[Inference API] Deduplicate PersistedConfig test utils class (elastic…

…#109862)
leemthompo · Jun 21, 2024 · 47ad822 · 47ad822
1 parent 57f4870
commit 47ad822
Show file tree

Hide file tree

Showing 8 changed files with 29 additions and 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
@@ -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<String, Object> config, Map<String, Object> secrets) {}
 }
diff --git a/...a/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioServiceTests.java b/...a/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<String, Object> getRequestConfigMap(
         );
     }
 
-    private record PeristedConfigRecord(Map<String, Object> config, Map<String, Object> secrets) {}
-
-    private PeristedConfigRecord getPersistedConfigMap(
+    private PersistedConfig getPersistedConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,
         Map<String, Object> 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<String, Object> serviceSettings, Map<String, Object> taskSettings) {
-
-        return new PeristedConfigRecord(
-            new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)),
-            null
-        );
-    }
-
     private static Map<String, Object> getEmbeddingsServiceSettingsMap(
         String target,
         String provider,

diff --git a/.../java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiServiceTests.java b/.../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<String, Object> getRequestConfigMap(
         );
     }
 
-    private PeristedConfig getPersistedConfigMap(
+    private PersistedConfig getPersistedConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,
         Map<String, Object> 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<String, Object> serviceSettings, Map<String, Object> taskSettings) {
-
-        return new PeristedConfig(
+    private PersistedConfig getPersistedConfigMap(Map<String, Object> serviceSettings, Map<String, Object> taskSettings) {
+        return new PersistedConfig(
             new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)),
             null
         );
     }
-
-    private record PeristedConfig(Map<String, Object> config, Map<String, Object> secrets) {}
 }
diff --git a/...e/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java b/...e/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<String, Object> serviceSetting
             null
         );
     }
-
-    private record PersistedConfig(Map<String, Object> config, Map<String, Object> secrets) {}
 }
diff --git a/...org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioServiceTests.java b/...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<String, Object> serviceSetting
             null
         );
     }
-
-    private record PersistedConfig(Map<String, Object> config, Map<String, Object> secrets) {}
-
 }
diff --git a/.../java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java b/.../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,30 +712,25 @@ private Map<String, Object> getRequestConfigMap(Map<String, Object> serviceSetti
         return new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, builtServiceSettings));
     }
 
-    private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap(Map<String, Object> serviceSettings) {
+    private PersistedConfig getPersistedConfigMap(Map<String, Object> serviceSettings) {
         return getPersistedConfigMap(serviceSettings, Map.of(), null);
     }
 
-    private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap(
-        Map<String, Object> serviceSettings,
-        @Nullable Map<String, Object> secretSettings
-    ) {
+    private PersistedConfig getPersistedConfigMap(Map<String, Object> serviceSettings, @Nullable Map<String, Object> secretSettings) {
         return getPersistedConfigMap(serviceSettings, Map.of(), secretSettings);
     }
 
-    private HuggingFaceServiceTests.PeristedConfig getPersistedConfigMap(
+    private PersistedConfig getPersistedConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,
         Map<String, Object> secretSettings
     ) {
 
         var secrets = secretSettings == null ? null : new HashMap<String, Object>(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<String, Object> config, Map<String, Object> secrets) {}
 }
diff --git a/...src/test/java/org/elasticsearch/xpack/inference/services/mistral/MistralServiceTests.java b/...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<String, Object> getRequestConfigMap(
         );
     }
 
-    private record PeristedConfigRecord(Map<String, Object> config, Map<String, Object> secrets) {}
-
-    private PeristedConfigRecord getPersistedConfigMap(
+    private PersistedConfig getPersistedConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,
         Map<String, Object> 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<String, Object> serviceSettings, Map<String, Object> taskSettings) {
-
-        return new PeristedConfigRecord(
-            new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)),
-            null
-        );
-    }
-
     private static Map<String, Object> getEmbeddingsServiceSettingsMap(
         String model,
         @Nullable Integer dimensions,

diff --git a/...e/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java b/...e/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<String, Object> getRequestConfigMap(
         );
     }
 
-    private PeristedConfig getPersistedConfigMap(
+    private PersistedConfig getPersistedConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,
         Map<String, Object> 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<String, Object> serviceSettings, Map<String, Object> taskSettings) {
+    private PersistedConfig getPersistedConfigMap(Map<String, Object> serviceSettings, Map<String, Object> taskSettings) {
 
-        return new PeristedConfig(
+        return new PersistedConfig(
             new HashMap<>(Map.of(ModelConfigurations.SERVICE_SETTINGS, serviceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)),
             null
         );
     }
-
-    private record PeristedConfig(Map<String, Object> config, Map<String, Object> secrets) {}
 }