fix test

Signed-off-by: Yaliang Wu <ylwu@amazon.com>
zane-neo · Jun 12, 2024 · 9c4df89 · 9c4df89
1 parent 38d8ddf
commit 9c4df89
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/...earch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java b/...earch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java
@@ -11,6 +11,7 @@
 import org.opensearch.ml.common.dataset.remote.RemoteInferenceInputDataSet;
 import org.opensearch.ml.common.input.MLInput;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestBedRockInferenceIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestBedRockInferenceIT.java
@@ -73,7 +73,7 @@ public void test_bedrock_multimodal_model() throws Exception {
 
             TextDocsInputDataSet inputDataSet = TextDocsInputDataSet.builder().docs(List.of("hello", imageBase64)).build();
             MLInput mlInput = MLInput.builder().inputDataset(inputDataSet).algorithm(FunctionName.TEXT_EMBEDDING).build();
-            Map inferenceResult = predictTextEmbedding(modelId, mlInput);
+            Map inferenceResult = predictTextEmbeddingModel(modelId, mlInput);
             assertTrue(errorMsg, inferenceResult.containsKey("inference_results"));
             List output = (List) inferenceResult.get("inference_results");
             assertEquals(errorMsg, 1, output.size());
@@ -86,7 +86,7 @@ public void test_bedrock_multimodal_model() throws Exception {
             assertEquals(errorMsg, 1024, ((List) ((Map<?, ?>) outputList.get(0)).get("data")).size());
         }
     }
-  
+
     public void test_bedrock_embedding_model() throws Exception {
         // Skip test if key is null
         if (AWS_ACCESS_KEY_ID == null || AWS_SECRET_ACCESS_KEY == null || AWS_SESSION_TOKEN == null) {