Remove hardcoded test schema

Signed-off-by: b4sjoo <[email protected]>
opensearch-project · Dec 20, 2024 · 9252146 · 9252146
1 parent 9ff2d5d
commit 9252146
Showing 1 changed file with 12 additions and 105 deletions.
diff --git a/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java b/plugin/src/test/java/org/opensearch/ml/utils/MLNodeUtilsTests.java
@@ -6,6 +6,9 @@
 package org.opensearch.ml.utils;
 
 import static java.util.Collections.emptyMap;
+import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE;
+import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE;
+import static org.opensearch.ml.common.utils.ModelInterfaceUtils.BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE;
 import static org.opensearch.ml.utils.TestHelper.ML_ROLE;
 
 import java.io.IOException;
@@ -68,139 +71,43 @@ public void testValidateSchema() throws IOException {
 
     @Test
     public void testValidateEmbeddingInputWithGeneralEmbeddingRemoteSchema() throws IOException {
-        String schema = "{\n"
-            + "    \"type\": \"object\",\n"
-            + "    \"properties\": {\n"
-            + "        \"parameters\": {\n"
-            + "            \"type\": \"object\",\n"
-            + "            \"properties\": {\n"
-            + "                \"texts\": {\n"
-            + "                    \"type\": \"array\",\n"
-            + "                    \"items\": {\n"
-            + "                        \"type\": \"string\"\n"
-            + "                    }\n"
-            + "                }\n"
-            + "            },\n"
-            + "            \"required\": [\n"
-            + "                \"texts\"\n"
-            + "            ]\n"
-            + "        }\n"
-            + "    }\n"
-            + "}";
+        String schema = BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE.get("input");
         String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
         MLNodeUtils.validateSchema(schema, json);
     }
 
     @Test
     public void testValidateRemoteInputWithGeneralEmbeddingRemoteSchema() throws IOException {
-        String schema = "{\n"
-            + "    \"type\": \"object\",\n"
-            + "    \"properties\": {\n"
-            + "        \"parameters\": {\n"
-            + "            \"type\": \"object\",\n"
-            + "            \"properties\": {\n"
-            + "                \"texts\": {\n"
-            + "                    \"type\": \"array\",\n"
-            + "                    \"items\": {\n"
-            + "                        \"type\": \"string\"\n"
-            + "                    }\n"
-            + "                }\n"
-            + "            },\n"
-            + "            \"required\": [\n"
-            + "                \"texts\"\n"
-            + "            ]\n"
-            + "        }\n"
-            + "    }\n"
-            + "}";
+        String schema = BEDROCK_COHERE_EMBED_ENGLISH_V3_MODEL_INTERFACE.get("input");
         String json = "{\"parameters\": {\"texts\": [\"Hello\",\"world\"]}}";
         MLNodeUtils.validateSchema(schema, json);
     }
 
     @Test
     public void testValidateEmbeddingInputWithTitanTextRemoteSchema() throws IOException {
-        String schema = "{\n"
-            + "    \"type\": \"object\",\n"
-            + "    \"properties\": {\n"
-            + "        \"parameters\": {\n"
-            + "            \"type\": \"object\",\n"
-            + "            \"properties\": {\n"
-            + "                \"inputText\": {\n"
-            + "                    \"type\": \"string\"\n"
-            + "                }\n"
-            + "            },\n"
-            + "            \"required\": [\n"
-            + "                \"inputText\"\n"
-            + "            ]\n"
-            + "        }\n"
-            + "    }\n"
-            + "}";
+        String schema = BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE.get("input");
         String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
         MLNodeUtils.validateSchema(schema, json);
     }
 
     @Test
     public void testValidateRemoteInputWithTitanTextRemoteSchema() throws IOException {
-        String schema = "{\n"
-            + "    \"type\": \"object\",\n"
-            + "    \"properties\": {\n"
-            + "        \"parameters\": {\n"
-            + "            \"type\": \"object\",\n"
-            + "            \"properties\": {\n"
-            + "                \"inputText\": {\n"
-            + "                    \"type\": \"string\"\n"
-            + "                }\n"
-            + "            },\n"
-            + "            \"required\": [\n"
-            + "                \"inputText\"\n"
-            + "            ]\n"
-            + "        }\n"
-            + "    }\n"
-            + "}";
+        String schema = BEDROCK_TITAN_EMBED_TEXT_V1_MODEL_INTERFACE.get("input");
         String json = "{\"parameters\": {\"inputText\": \"Say this is a test\"}}";
         MLNodeUtils.validateSchema(schema, json);
     }
 
     @Test
-    public void testValidateEmbeddingInputWithTitanImageRemoteSchema() throws IOException {
-        String schema = "{\n"
-            + "    \"type\": \"object\",\n"
-            + "    \"properties\": {\n"
-            + "        \"parameters\": {\n"
-            + "            \"type\": \"object\",\n"
-            + "            \"properties\": {\n"
-            + "                \"inputText\": {\n"
-            + "                    \"type\": \"string\"\n"
-            + "                },\n"
-            + "                \"inputImage\": {\n"
-            + "                    \"type\": \"string\"\n"
-            + "                }\n"
-            + "            }\n"
-            + "        }\n"
-            + "    }\n"
-            + "}";
+    public void testValidateEmbeddingInputWithTitanMultiModalRemoteSchema() throws IOException {
+        String schema = BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE.get("input");
         String json = "{\"text_docs\":[ \"today is sunny\", \"today is sunny\"]}";
         MLNodeUtils.validateSchema(schema, json);
     }
 
     @Test
-    public void testValidateRemoteInputWithTitanImageRemoteSchema() throws IOException {
-        String schema = "{\n"
-            + "    \"type\": \"object\",\n"
-            + "    \"properties\": {\n"
-            + "        \"parameters\": {\n"
-            + "            \"type\": \"object\",\n"
-            + "            \"properties\": {\n"
-            + "                \"inputText\": {\n"
-            + "                    \"type\": \"string\"\n"
-            + "                },\n"
-            + "                \"inputImage\": {\n"
-            + "                    \"type\": \"string\"\n"
-            + "                }\n"
-            + "            }\n"
-            + "        }\n"
-            + "    }\n"
-            + "}";
-        String json = "{{\n"
+    public void testValidateRemoteInputWithTitanMultiModalRemoteSchema() throws IOException {
+        String schema = BEDROCK_TITAN_EMBED_MULTI_MODAL_V1_MODEL_INTERFACE.get("input");
+        String json = "{\n"
             + "  \"parameters\": {\n"
             + "    \"inputText\": \"Say this is a test\",\n"
             + "    \"inputImage\": \"/9jk=\"\n"