From 47ad8224caf4a5f0908907145b614872aeeb5c0f Mon Sep 17 00:00:00 2001 From: Tim Grein Date: Fri, 21 Jun 2024 10:50:08 +0200 Subject: [PATCH] [Inference API] Deduplicate PersistedConfig test utils class (#109862) --- .../org/elasticsearch/xpack/inference/Utils.java | 2 ++ .../azureaistudio/AzureAiStudioServiceTests.java | 15 +++------------ .../azureopenai/AzureOpenAiServiceTests.java | 14 ++++++-------- .../services/cohere/CohereServiceTests.java | 5 ++--- .../GoogleAiStudioServiceTests.java | 4 +--- .../huggingface/HuggingFaceServiceTests.java | 16 ++++++---------- .../services/mistral/MistralServiceTests.java | 15 +++------------ .../services/openai/OpenAiServiceTests.java | 13 ++++++------- 8 files changed, 29 insertions(+), 55 deletions(-) diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java index 58603526a9c56..a352116278e7a 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java @@ -30,6 +30,7 @@ import java.util.Collection; import java.util.List; +import java.util.Map; import java.util.concurrent.CountDownLatch; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -159,4 +160,5 @@ public static SimilarityMeasure randomSimilarityMeasure() { return randomFrom(SimilarityMeasure.values()); } + public record PersistedConfig(Map config, Map secrets) {} } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java index 709cc4d3494fd..d26b02ddba62b 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java @@ -60,6 +60,7 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.getInvalidModel; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; @@ -1068,28 +1069,18 @@ private Map getRequestConfigMap( ); } - private record PeristedConfigRecord(Map config, Map secrets) {} - - private PeristedConfigRecord getPersistedConfigMap( + private PersistedConfig getPersistedConfigMap( Map serviceSettings, Map taskSettings, Map secretSettings ) { - return new PeristedConfigRecord( + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), new HashMap<>(Map.of(ModelSecrets.SECRET_SETTINGS, secretSettings)) ); } - private PeristedConfigRecord getPersistedConfigMap(Map serviceSettings, Map taskSettings) { - - return new PeristedConfigRecord( - new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), - null - ); - } - private static Map getEmbeddingsServiceSettingsMap( String target, String provider, diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java index de474ea1b4237..c3e8eb5c621d2 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java @@ -54,6 +54,7 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.getInvalidModel; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; @@ -394,7 +395,7 @@ public void testParsePersistedConfigWithSecrets_NotThrowWhenAnExtraKeyExistsInSe getAzureOpenAiRequestTaskSettingsMap("user"), getAzureOpenAiSecretSettingsMap("secret", null) ); - persistedConfig.secrets.put("extra_key", "value"); + persistedConfig.secrets().put("extra_key", "value"); var model = service.parsePersistedConfigWithSecrets( "id", @@ -1159,25 +1160,22 @@ private Map getRequestConfigMap( ); } - private PeristedConfig getPersistedConfigMap( + private PersistedConfig getPersistedConfigMap( Map serviceSettings, Map taskSettings, Map secretSettings ) { - return new PeristedConfig( + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), new HashMap<>(Map.of(ModelSecrets.SECRET_SETTINGS, secretSettings)) ); } - private PeristedConfig getPersistedConfigMap(Map serviceSettings, Map taskSettings) { - - return new PeristedConfig( + private PersistedConfig getPersistedConfigMap(Map serviceSettings, Map taskSettings) { + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), null ); } - - private record PeristedConfig(Map config, Map secrets) {} } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java index 5b3cb9eade9de..e28ca71c30ff8 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java @@ -58,6 +58,7 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.getInvalidModel; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; @@ -421,7 +422,7 @@ public void testParsePersistedConfigWithSecrets_NotThrowWhenAnExtraKeyExistsInSe getTaskSettingsMap(null, null), getSecretSettingsMap("secret") ); - persistedConfig.secrets.put("extra_key", "value"); + persistedConfig.secrets().put("extra_key", "value"); var model = service.parsePersistedConfigWithSecrets( "id", @@ -1394,6 +1395,4 @@ private PersistedConfig getPersistedConfigMap(Map serviceSetting null ); } - - private record PersistedConfig(Map config, Map secrets) {} } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java index 1cdd7997b96c0..45dd8ad7b33bd 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java @@ -57,6 +57,7 @@ import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.getInvalidModel; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; @@ -963,7 +964,4 @@ private PersistedConfig getPersistedConfigMap(Map serviceSetting null ); } - - private record PersistedConfig(Map config, Map secrets) {} - } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java index a855437ce0738..de5c7ec83d57e 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java @@ -56,6 +56,7 @@ import java.util.concurrent.TimeUnit; import static org.elasticsearch.xpack.core.inference.results.InferenceChunkedTextEmbeddingFloatResultsTests.asMapWithListsInsteadOfArrays; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap; @@ -278,7 +279,7 @@ public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExists public void testParsePersistedConfigWithSecrets_DoesNotThrowWhenAnExtraKeyExistsInSecrets() throws IOException { try (var service = createHuggingFaceService()) { var persistedConfig = getPersistedConfigMap(getServiceSettingsMap("url"), getSecretSettingsMap("secret")); - persistedConfig.secrets.put("extra_key", "value"); + persistedConfig.secrets().put("extra_key", "value"); var model = service.parsePersistedConfigWithSecrets( "id", @@ -711,18 +712,15 @@ private Map getRequestConfigMap(Map serviceSetti return new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, builtServiceSettings)); } - private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap(Map serviceSettings) { + private PersistedConfig getPersistedConfigMap(Map serviceSettings) { return getPersistedConfigMap(serviceSettings, Map.of(), null); } - private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap( - Map serviceSettings, - @Nullable Map secretSettings - ) { + private PersistedConfig getPersistedConfigMap(Map serviceSettings, @Nullable Map secretSettings) { return getPersistedConfigMap(serviceSettings, Map.of(), secretSettings); } - private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap( + private PersistedConfig getPersistedConfigMap( Map serviceSettings, Map taskSettings, Map secretSettings @@ -730,11 +728,9 @@ private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap( var secrets = secretSettings == null ? null : new HashMap(Map.of(ModelSecrets.SECRET_SETTINGS, secretSettings)); - return new HuggingFaceServiceTests.PeristedConfig( + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), secrets ); } - - private record PeristedConfig(Map config, Map secrets) {} } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java index 508d5a97fe564..ba37203d9e5d6 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java @@ -54,6 +54,7 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.getInvalidModel; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; @@ -586,28 +587,18 @@ private Map getRequestConfigMap( ); } - private record PeristedConfigRecord(Map config, Map secrets) {} - - private PeristedConfigRecord getPersistedConfigMap( + private PersistedConfig getPersistedConfigMap( Map serviceSettings, Map taskSettings, Map secretSettings ) { - return new PeristedConfigRecord( + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), new HashMap<>(Map.of(ModelSecrets.SECRET_SETTINGS, secretSettings)) ); } - private PeristedConfigRecord getPersistedConfigMap(Map serviceSettings, Map taskSettings) { - - return new PeristedConfigRecord( - new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), - null - ); - } - private static Map getEmbeddingsServiceSettingsMap( String model, @Nullable Integer dimensions, diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java index e0e1ee3e81aef..2fc049dd3a5f6 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java @@ -55,6 +55,7 @@ import java.util.Set; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.xpack.inference.Utils.PersistedConfig; import static org.elasticsearch.xpack.inference.Utils.getInvalidModel; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; @@ -480,7 +481,7 @@ public void testParsePersistedConfigWithSecrets_NotThrowWhenAnExtraKeyExistsInSe getTaskSettingsMap("user"), getSecretSettingsMap("secret") ); - persistedConfig.secrets.put("extra_key", "value"); + persistedConfig.secrets().put("extra_key", "value"); var model = service.parsePersistedConfigWithSecrets( "id", @@ -1308,25 +1309,23 @@ private Map getRequestConfigMap( ); } - private PeristedConfig getPersistedConfigMap( + private PersistedConfig getPersistedConfigMap( Map serviceSettings, Map taskSettings, Map secretSettings ) { - return new PeristedConfig( + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), new HashMap<>(Map.of(ModelSecrets.SECRET_SETTINGS, secretSettings)) ); } - private PeristedConfig getPersistedConfigMap(Map serviceSettings, Map taskSettings) { + private PersistedConfig getPersistedConfigMap(Map serviceSettings, Map taskSettings) { - return new PeristedConfig( + return new PersistedConfig( new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)), null ); } - - private record PeristedConfig(Map config, Map secrets) {} }