From 9252146d26b488ece78a5514f0d56fb884a988ba Mon Sep 17 00:00:00 2001 From: b4sjoo Date: Fri, 20 Dec 2024 19:29:24 +0000 Subject: [PATCH] Remove hardcoded test schema Signed-off-by: b4sjoo --- .../opensearch/ml/utils/MLNodeUtilsTests.java | 117 ++---------------- 1 file changed, 12 insertions(+), 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 index c0f6d35ee7..40bb230bbf 100644 --- 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"