From 1b15536b6ff4e47af617f7462aaac50c500575cc Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Tue, 1 Jul 2025 18:58:00 -0400 Subject: [PATCH 1/5] Added SemanticTextIndexOptionsIT --- .../SemanticTextIndexOptionsIT.java | 232 ++++++++++++++++++ 1 file changed, 232 insertions(+) create mode 100644 x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java new file mode 100644 index 0000000000000..36e74f047fb99 --- /dev/null +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java @@ -0,0 +1,232 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.integration; + +import org.elasticsearch.action.admin.indices.mapping.get.GetFieldMappingsAction; +import org.elasticsearch.action.admin.indices.mapping.get.GetFieldMappingsRequest; +import org.elasticsearch.action.support.IndicesOptions; +import org.elasticsearch.common.bytes.BytesReference; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.common.xcontent.XContentHelper; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; +import org.elasticsearch.index.mapper.vectors.IndexOptions; +import org.elasticsearch.inference.TaskType; +import org.elasticsearch.license.License; +import org.elasticsearch.license.LicenseSettings; +import org.elasticsearch.license.PostStartBasicAction; +import org.elasticsearch.license.PostStartBasicRequest; +import org.elasticsearch.license.PutLicenseAction; +import org.elasticsearch.license.PutLicenseRequest; +import org.elasticsearch.license.TestUtils; +import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.reindex.ReindexPlugin; +import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.xcontent.ToXContent; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentFactory; +import org.elasticsearch.xcontent.XContentType; +import org.elasticsearch.xpack.core.inference.action.DeleteInferenceEndpointAction; +import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction; +import org.elasticsearch.xpack.inference.LocalStateInferencePlugin; +import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; +import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension; +import org.elasticsearch.xpack.inference.mock.TestInferenceServicePlugin; +import org.elasticsearch.xpack.inference.mock.TestSparseInferenceServiceExtension; +import org.junit.After; +import org.junit.Before; + +import java.io.IOException; +import java.util.Collection; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked; +import static org.hamcrest.CoreMatchers.equalTo; + +public class SemanticTextIndexOptionsIT extends ESIntegTestCase { + private static final String INDEX_NAME = "test-index"; + private static final Map BBQ_COMPATIBLE_SERVICE_SETTINGS = Map.of( + "model", + "my_model", + "dimensions", + 256, + "similarity", + "cosine", + "api_key", + "my_api_key" + ); + + private final Map inferenceIds = new HashMap<>(); + + @Override + protected Settings nodeSettings(int nodeOrdinal, Settings otherSettings) { + return Settings.builder().put(LicenseSettings.SELF_GENERATED_LICENSE_TYPE.getKey(), "trial").build(); + } + + @Override + protected Collection> nodePlugins() { + return List.of(LocalStateInferencePlugin.class, TestInferenceServicePlugin.class, ReindexPlugin.class); + } + + @Before + public void resetLicense() throws Exception { + setLicense("trial"); + } + + @After + public void cleanUp() { + assertAcked( + safeGet( + client().admin() + .indices() + .prepareDelete(INDEX_NAME) + .setIndicesOptions( + IndicesOptions.builder().concreteTargetOptions(new IndicesOptions.ConcreteTargetOptions(true)).build() + ) + .execute() + ) + ); + + for (var entry : inferenceIds.entrySet()) { + assertAcked( + safeGet( + client().execute( + DeleteInferenceEndpointAction.INSTANCE, + new DeleteInferenceEndpointAction.Request(entry.getKey(), entry.getValue(), true, false) + ) + ) + ); + } + } + + public void testValidateIndexOptionsWithBasicLicense() throws Exception { + final String inferenceId = "test-inference-id-1"; + final String inferenceFieldName = "inference_field"; + createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); + + setLicense("basic"); + IndexOptions indexOptions = new DenseVectorFieldMapper.Int8HnswIndexOptions( + randomIntBetween(1, 100), + randomIntBetween(1, 10_000), + null, + null + ); + assertAcked( + safeGet(prepareCreate(INDEX_NAME).setMapping(generateMapping(inferenceFieldName, inferenceId, indexOptions)).execute()) + ); + + final Map expectedFieldMapping = generateExpectedFieldMapping(inferenceId, inferenceFieldName, indexOptions); + var getFieldMappingsResponse = safeGet( + client().execute(GetFieldMappingsAction.INSTANCE, new GetFieldMappingsRequest().indices(INDEX_NAME).fields(inferenceFieldName)) + ); + assertThat(getFieldMappingsResponse.fieldMappings(INDEX_NAME, inferenceFieldName).sourceAsMap(), equalTo(expectedFieldMapping)); + } + + private void createInferenceEndpoint(TaskType taskType, String inferenceId, Map serviceSettings) throws IOException { + final String service = switch (taskType) { + case TEXT_EMBEDDING -> TestDenseInferenceServiceExtension.TestInferenceService.NAME; + case SPARSE_EMBEDDING -> TestSparseInferenceServiceExtension.TestInferenceService.NAME; + default -> throw new IllegalArgumentException("Unhandled task type [" + taskType + "]"); + }; + + final BytesReference content; + try (XContentBuilder builder = XContentFactory.jsonBuilder()) { + builder.startObject(); + builder.field("service", service); + builder.field("service_settings", serviceSettings); + builder.endObject(); + + content = BytesReference.bytes(builder); + } + + PutInferenceModelAction.Request request = new PutInferenceModelAction.Request( + taskType, + inferenceId, + content, + XContentType.JSON, + TEST_REQUEST_TIMEOUT + ); + var responseFuture = client().execute(PutInferenceModelAction.INSTANCE, request); + assertThat(responseFuture.actionGet(TEST_REQUEST_TIMEOUT).getModel().getInferenceEntityId(), equalTo(inferenceId)); + + inferenceIds.put(inferenceId, taskType); + } + + private static XContentBuilder generateMapping(String inferenceFieldName, String inferenceId, @Nullable IndexOptions indexOptions) + throws IOException { + XContentBuilder mapping = XContentFactory.jsonBuilder(); + mapping.startObject(); + mapping.field("properties"); + generateFieldMapping(mapping, inferenceFieldName, inferenceId, indexOptions); + mapping.endObject(); + + return mapping; + } + + private static void generateFieldMapping( + XContentBuilder builder, + String inferenceFieldName, + String inferenceId, + @Nullable IndexOptions indexOptions + ) throws IOException { + builder.startObject(); + builder.startObject(inferenceFieldName); + builder.field("type", SemanticTextFieldMapper.CONTENT_TYPE); + builder.field("inference_id", inferenceId); + if (indexOptions != null) { + builder.startObject("index_options"); + if (indexOptions instanceof DenseVectorFieldMapper.DenseVectorIndexOptions) { + builder.field("dense_vector"); + indexOptions.toXContent(builder, ToXContent.EMPTY_PARAMS); + } + builder.endObject(); + } + builder.endObject(); + builder.endObject(); + } + + private static Map generateExpectedFieldMapping( + String inferenceFieldName, + String inferenceId, + @Nullable IndexOptions indexOptions + ) throws IOException { + Map expectedFieldMapping; + try (XContentBuilder builder = XContentFactory.jsonBuilder()) { + generateFieldMapping(builder, inferenceFieldName, inferenceId, indexOptions); + expectedFieldMapping = XContentHelper.convertToMap(BytesReference.bytes(builder), false, XContentType.JSON).v2(); + } + + return expectedFieldMapping; + } + + private static void setLicense(String type) throws Exception { + if (type.equals("basic")) { + assertAcked( + safeGet( + client().execute( + PostStartBasicAction.INSTANCE, + new PostStartBasicRequest(TEST_REQUEST_TIMEOUT, TEST_REQUEST_TIMEOUT).acknowledge(true) + ) + ) + ); + } else { + License license = TestUtils.generateSignedLicense(type, License.VERSION_CURRENT, -1, TimeValue.timeValueHours(24)); + assertAcked( + safeGet( + client().execute( + PutLicenseAction.INSTANCE, + new PutLicenseRequest(TEST_REQUEST_TIMEOUT, TEST_REQUEST_TIMEOUT).license(license) + ) + ) + ); + } + } +} From bae2d145ac9114924bcb783cf6975f9a3ac4e063 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Tue, 1 Jul 2025 19:23:48 -0400 Subject: [PATCH 2/5] Added default BBQ index options test --- .../SemanticTextIndexOptionsIT.java | 50 ++++++++++++++++++- .../mapper/SemanticTextFieldMapper.java | 2 +- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java index 36e74f047fb99..f4d2c81c0c326 100644 --- a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java @@ -123,13 +123,40 @@ public void testValidateIndexOptionsWithBasicLicense() throws Exception { safeGet(prepareCreate(INDEX_NAME).setMapping(generateMapping(inferenceFieldName, inferenceId, indexOptions)).execute()) ); - final Map expectedFieldMapping = generateExpectedFieldMapping(inferenceId, inferenceFieldName, indexOptions); + final Map expectedFieldMapping = generateExpectedFieldMapping(inferenceFieldName, inferenceId, indexOptions); var getFieldMappingsResponse = safeGet( client().execute(GetFieldMappingsAction.INSTANCE, new GetFieldMappingsRequest().indices(INDEX_NAME).fields(inferenceFieldName)) ); assertThat(getFieldMappingsResponse.fieldMappings(INDEX_NAME, inferenceFieldName).sourceAsMap(), equalTo(expectedFieldMapping)); } + public void testSetDefaultBBQIndexOptionsWithBasicLicense() throws Exception { + final String inferenceId = "test-inference-id-2"; + final String inferenceFieldName = "inference_field"; + createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); + + setLicense("basic"); + assertAcked(safeGet(prepareCreate(INDEX_NAME).setMapping(generateMapping(inferenceFieldName, inferenceId, null)).execute())); + + final Map expectedFieldMapping = generateExpectedFieldMapping( + inferenceFieldName, + inferenceId, + SemanticTextFieldMapper.defaultBbqHnswDenseVectorIndexOptions() + ); + var getFieldMappingsResponse = safeGet( + client().execute( + GetFieldMappingsAction.INSTANCE, + new GetFieldMappingsRequest().indices(INDEX_NAME).fields(inferenceFieldName).includeDefaults(true) + ) + ); + + // Filter out null/empty values from params we didn't set to make comparison easier + Map actualFieldMappings = filterNullOrEmptyValues( + getFieldMappingsResponse.fieldMappings(INDEX_NAME, inferenceFieldName).sourceAsMap() + ); + assertThat(actualFieldMappings, equalTo(expectedFieldMapping)); + } + private void createInferenceEndpoint(TaskType taskType, String inferenceId, Map serviceSettings) throws IOException { final String service = switch (taskType) { case TEXT_EMBEDDING -> TestDenseInferenceServiceExtension.TestInferenceService.NAME; @@ -207,6 +234,27 @@ private static Map generateExpectedFieldMapping( return expectedFieldMapping; } + @SuppressWarnings("unchecked") + private static Map filterNullOrEmptyValues(Map map) { + Map filteredMap = new HashMap<>(); + for (var entry : map.entrySet()) { + Object value = entry.getValue(); + if (entry.getValue() instanceof Map mapValue) { + if (mapValue.isEmpty()) { + continue; + } + + value = filterNullOrEmptyValues((Map) mapValue); + } + + if (value != null) { + filteredMap.put(entry.getKey(), value); + } + } + + return filteredMap; + } + private static void setLicense(String type) throws Exception { if (type.equals("basic")) { assertAcked( diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java index fd5f1ce2735a9..31e97098e2877 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java @@ -1250,7 +1250,7 @@ static boolean indexVersionDefaultsToBbqHnsw(IndexVersion indexVersion) { || indexVersion.between(SEMANTIC_TEXT_DEFAULTS_TO_BBQ_BACKPORT_8_X, IndexVersions.UPGRADE_TO_LUCENE_10_0_0); } - static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() { + public static DenseVectorFieldMapper.DenseVectorIndexOptions defaultBbqHnswDenseVectorIndexOptions() { int m = Lucene99HnswVectorsFormat.DEFAULT_MAX_CONN; int efConstruction = Lucene99HnswVectorsFormat.DEFAULT_BEAM_WIDTH; DenseVectorFieldMapper.RescoreVector rescoreVector = new DenseVectorFieldMapper.RescoreVector(DEFAULT_RESCORE_OVERSAMPLE); From 24d628f6b2ffdd1dfaf135cb5bb74997f9d69166 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Tue, 1 Jul 2025 19:39:17 -0400 Subject: [PATCH 3/5] Restart cluster after downgrading license --- .../SemanticTextIndexOptionsIT.java | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java index f4d2c81c0c326..960c416349112 100644 --- a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java @@ -18,6 +18,7 @@ import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper; import org.elasticsearch.index.mapper.vectors.IndexOptions; import org.elasticsearch.inference.TaskType; +import org.elasticsearch.license.GetLicenseAction; import org.elasticsearch.license.License; import org.elasticsearch.license.LicenseSettings; import org.elasticsearch.license.PostStartBasicAction; @@ -26,14 +27,17 @@ import org.elasticsearch.license.PutLicenseRequest; import org.elasticsearch.license.TestUtils; import org.elasticsearch.plugins.Plugin; +import org.elasticsearch.protocol.xpack.license.GetLicenseRequest; import org.elasticsearch.reindex.ReindexPlugin; import org.elasticsearch.test.ESIntegTestCase; +import org.elasticsearch.test.InternalTestCluster; import org.elasticsearch.xcontent.ToXContent; import org.elasticsearch.xcontent.XContentBuilder; import org.elasticsearch.xcontent.XContentFactory; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.core.inference.action.DeleteInferenceEndpointAction; import org.elasticsearch.xpack.core.inference.action.PutInferenceModelAction; +import org.elasticsearch.xpack.inference.InferenceIndex; import org.elasticsearch.xpack.inference.LocalStateInferencePlugin; import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper; import org.elasticsearch.xpack.inference.mock.TestDenseInferenceServiceExtension; @@ -112,7 +116,12 @@ public void testValidateIndexOptionsWithBasicLicense() throws Exception { final String inferenceFieldName = "inference_field"; createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); + // Downgrade the license and restart the cluster to force the model registry to rebuild setLicense("basic"); + internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); + ensureGreen(InferenceIndex.INDEX_NAME); + assertLicense(License.LicenseType.BASIC); + IndexOptions indexOptions = new DenseVectorFieldMapper.Int8HnswIndexOptions( randomIntBetween(1, 100), randomIntBetween(1, 10_000), @@ -135,7 +144,12 @@ public void testSetDefaultBBQIndexOptionsWithBasicLicense() throws Exception { final String inferenceFieldName = "inference_field"; createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); + // Downgrade the license and restart the cluster to force the model registry to rebuild setLicense("basic"); + internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); + ensureGreen(InferenceIndex.INDEX_NAME); + assertLicense(License.LicenseType.BASIC); + assertAcked(safeGet(prepareCreate(INDEX_NAME).setMapping(generateMapping(inferenceFieldName, inferenceId, null)).execute())); final Map expectedFieldMapping = generateExpectedFieldMapping( @@ -255,6 +269,7 @@ private static Map filterNullOrEmptyValues(Map m return filteredMap; } + // TODO: Pass LicenseType to this method private static void setLicense(String type) throws Exception { if (type.equals("basic")) { assertAcked( @@ -277,4 +292,9 @@ private static void setLicense(String type) throws Exception { ); } } + + private static void assertLicense(License.LicenseType type) { + var getLicenseResponse = safeGet(client().execute(GetLicenseAction.INSTANCE, new GetLicenseRequest(TEST_REQUEST_TIMEOUT))); + assertThat(getLicenseResponse.license().type(), equalTo(type.getTypeName())); + } } From 6bb6075b12fb7f6c94f6436f9360e5ef74f57483 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Wed, 2 Jul 2025 08:46:10 -0400 Subject: [PATCH 4/5] Pass LicenseType to setLicense --- .../SemanticTextIndexOptionsIT.java | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java index 960c416349112..c814fb0593860 100644 --- a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java @@ -82,7 +82,7 @@ protected Collection> nodePlugins() { @Before public void resetLicense() throws Exception { - setLicense("trial"); + setLicense(License.LicenseType.TRIAL); } @After @@ -117,7 +117,7 @@ public void testValidateIndexOptionsWithBasicLicense() throws Exception { createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); // Downgrade the license and restart the cluster to force the model registry to rebuild - setLicense("basic"); + setLicense(License.LicenseType.BASIC); internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); ensureGreen(InferenceIndex.INDEX_NAME); assertLicense(License.LicenseType.BASIC); @@ -145,7 +145,7 @@ public void testSetDefaultBBQIndexOptionsWithBasicLicense() throws Exception { createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); // Downgrade the license and restart the cluster to force the model registry to rebuild - setLicense("basic"); + setLicense(License.LicenseType.BASIC); internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); ensureGreen(InferenceIndex.INDEX_NAME); assertLicense(License.LicenseType.BASIC); @@ -269,9 +269,8 @@ private static Map filterNullOrEmptyValues(Map m return filteredMap; } - // TODO: Pass LicenseType to this method - private static void setLicense(String type) throws Exception { - if (type.equals("basic")) { + private static void setLicense(License.LicenseType type) throws Exception { + if (type == License.LicenseType.BASIC) { assertAcked( safeGet( client().execute( @@ -281,7 +280,12 @@ private static void setLicense(String type) throws Exception { ) ); } else { - License license = TestUtils.generateSignedLicense(type, License.VERSION_CURRENT, -1, TimeValue.timeValueHours(24)); + License license = TestUtils.generateSignedLicense( + type.getTypeName(), + License.VERSION_CURRENT, + -1, + TimeValue.timeValueHours(24) + ); assertAcked( safeGet( client().execute( From d8ee685b6177ffea36c5a9df9050385e4a287993 Mon Sep 17 00:00:00 2001 From: Mike Pellegrini Date: Thu, 3 Jul 2025 13:18:24 -0400 Subject: [PATCH 5/5] Factor out duplicated code --- .../SemanticTextIndexOptionsIT.java | 42 ++++++++----------- 1 file changed, 17 insertions(+), 25 deletions(-) diff --git a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java index c814fb0593860..035ec4dc9593d 100644 --- a/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java +++ b/x-pack/plugin/inference/src/internalClusterTest/java/org/elasticsearch/xpack/inference/integration/SemanticTextIndexOptionsIT.java @@ -115,12 +115,7 @@ public void testValidateIndexOptionsWithBasicLicense() throws Exception { final String inferenceId = "test-inference-id-1"; final String inferenceFieldName = "inference_field"; createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); - - // Downgrade the license and restart the cluster to force the model registry to rebuild - setLicense(License.LicenseType.BASIC); - internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); - ensureGreen(InferenceIndex.INDEX_NAME); - assertLicense(License.LicenseType.BASIC); + downgradeLicenseAndRestartCluster(); IndexOptions indexOptions = new DenseVectorFieldMapper.Int8HnswIndexOptions( randomIntBetween(1, 100), @@ -133,22 +128,14 @@ public void testValidateIndexOptionsWithBasicLicense() throws Exception { ); final Map expectedFieldMapping = generateExpectedFieldMapping(inferenceFieldName, inferenceId, indexOptions); - var getFieldMappingsResponse = safeGet( - client().execute(GetFieldMappingsAction.INSTANCE, new GetFieldMappingsRequest().indices(INDEX_NAME).fields(inferenceFieldName)) - ); - assertThat(getFieldMappingsResponse.fieldMappings(INDEX_NAME, inferenceFieldName).sourceAsMap(), equalTo(expectedFieldMapping)); + assertThat(getFieldMappings(inferenceFieldName, false), equalTo(expectedFieldMapping)); } public void testSetDefaultBBQIndexOptionsWithBasicLicense() throws Exception { final String inferenceId = "test-inference-id-2"; final String inferenceFieldName = "inference_field"; createInferenceEndpoint(TaskType.TEXT_EMBEDDING, inferenceId, BBQ_COMPATIBLE_SERVICE_SETTINGS); - - // Downgrade the license and restart the cluster to force the model registry to rebuild - setLicense(License.LicenseType.BASIC); - internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); - ensureGreen(InferenceIndex.INDEX_NAME); - assertLicense(License.LicenseType.BASIC); + downgradeLicenseAndRestartCluster(); assertAcked(safeGet(prepareCreate(INDEX_NAME).setMapping(generateMapping(inferenceFieldName, inferenceId, null)).execute())); @@ -157,17 +144,9 @@ public void testSetDefaultBBQIndexOptionsWithBasicLicense() throws Exception { inferenceId, SemanticTextFieldMapper.defaultBbqHnswDenseVectorIndexOptions() ); - var getFieldMappingsResponse = safeGet( - client().execute( - GetFieldMappingsAction.INSTANCE, - new GetFieldMappingsRequest().indices(INDEX_NAME).fields(inferenceFieldName).includeDefaults(true) - ) - ); // Filter out null/empty values from params we didn't set to make comparison easier - Map actualFieldMappings = filterNullOrEmptyValues( - getFieldMappingsResponse.fieldMappings(INDEX_NAME, inferenceFieldName).sourceAsMap() - ); + Map actualFieldMappings = filterNullOrEmptyValues(getFieldMappings(inferenceFieldName, true)); assertThat(actualFieldMappings, equalTo(expectedFieldMapping)); } @@ -269,6 +248,11 @@ private static Map filterNullOrEmptyValues(Map m return filteredMap; } + private static Map getFieldMappings(String fieldName, boolean includeDefaults) { + var request = new GetFieldMappingsRequest().indices(INDEX_NAME).fields(fieldName).includeDefaults(includeDefaults); + return safeGet(client().execute(GetFieldMappingsAction.INSTANCE, request)).fieldMappings(INDEX_NAME, fieldName).sourceAsMap(); + } + private static void setLicense(License.LicenseType type) throws Exception { if (type == License.LicenseType.BASIC) { assertAcked( @@ -301,4 +285,12 @@ private static void assertLicense(License.LicenseType type) { var getLicenseResponse = safeGet(client().execute(GetLicenseAction.INSTANCE, new GetLicenseRequest(TEST_REQUEST_TIMEOUT))); assertThat(getLicenseResponse.license().type(), equalTo(type.getTypeName())); } + + private void downgradeLicenseAndRestartCluster() throws Exception { + // Downgrade the license and restart the cluster to force the model registry to rebuild + setLicense(License.LicenseType.BASIC); + internalCluster().fullRestart(new InternalTestCluster.RestartCallback()); + ensureGreen(InferenceIndex.INDEX_NAME); + assertLicense(License.LicenseType.BASIC); + } }