fix flaky test (opensearch-project#3598)

Signed-off-by: Mingshi Liu <mingshl@amazon.com>

Author: mingshl

plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceSearchRequestProcessorIT.java

@@ -170,21 +170,27 @@ public void setup() throws Exception {
             + "    \"properties\": {\n"
             + "      \"diary_embedding_size\": {\n"
             + "        \"type\": \"keyword\"\n"
+            + "      },\n"
+            + "      \"diary_embedding_size_int\": {\n"
+            + "        \"type\": \"integer\"\n"
             + "      }\n"
             + "    }\n"
             + "  }\n"
             + "}";
+
         String uploadDocumentRequestBodyDoc1 = "{\n"
             + "  \"id\": 1,\n"
             + "  \"diary\": [\"happy\",\"first day at school\"],\n"
             + "  \"diary_embedding_size\": \"1536\",\n" // embedding size for ada model
+            + "  \"diary_embedding_size_int\": 1536,\n"
             + "  \"weather\": \"rainy\"\n"
             + "  }";
 
         String uploadDocumentRequestBodyDoc2 = "{\n"
             + "  \"id\": 2,\n"
             + "  \"diary\": [\"bored\",\"at home\"],\n"
             + "  \"diary_embedding_size\": \"768\",\n"  // embedding size for local text embedding model
+            + "  \"diary_embedding_size_int\": 768,\n"
             + "  \"weather\": \"sunny\"\n"
             + "  }";
 
@@ -389,7 +395,7 @@ public void testMLInferenceProcessorRemoteModelOptionalInputs() throws Exception
             + "        \"model_id\": \""
             + this.bedrockMultiModalEmbeddingModelId
             + "\",\n"
-            + "        \"query_template\": \"{\\\"size\\\": 2,\\\"query\\\": {\\\"range\\\": {\\\"diary_embedding_size\\\": {\\\"gte\\\": ${modelPrediction}}}}}\",\n"
+            + "        \"query_template\": \"{\\\"size\\\": 2,\\\"query\\\": {\\\"range\\\": {\\\"diary_embedding_size_int\\\": {\\\"gte\\\": ${modelPrediction}}}}}\",\n"
             + "        \"optional_input_map\": [\n"
             + "          {\n"
             + "            \"inputText\": \"query.term.diary.value\",\n"
@@ -415,9 +421,8 @@ public void testMLInferenceProcessorRemoteModelOptionalInputs() throws Exception
         createSearchPipelineProcessor(createPipelineRequestBody, pipelineName);
 
         Map response = searchWithPipeline(client(), index_name, pipelineName, query);
-
         assertEquals((int) JsonPath.parse(response).read("$.hits.hits.length()"), 1);
-        Assert.assertEquals(JsonPath.parse(response).read("$.hits.hits[0]._source.diary_embedding_size"), "1536");
+        assertEquals((double) JsonPath.parse(response).read("$.hits.hits[0]._source.diary_embedding_size_int"), 1536.0, 0.0001);
     }
 
     /**

plugin/src/test/java/org/opensearch/ml/rest/RestMLInferenceSearchResponseProcessorIT.java

@@ -4,7 +4,6 @@
  */
 package org.opensearch.ml.rest;
 
-import static org.junit.Assert.assertEquals;
 import static org.opensearch.ml.common.MLModel.MODEL_ID_FIELD;
 import static org.opensearch.ml.utils.TestData.SENTENCE_TRANSFORMER_MODEL_URL;
 import static org.opensearch.ml.utils.TestHelper.makeRequest;
@@ -439,7 +438,7 @@ public void testMLInferenceProcessorRemoteModelOptionalInputs() throws Exception
             return;
         }
         String createPipelineRequestBody = "{\n"
-            + "  \"response\": [\n"
+            + "  \"response_processors\": [\n"
             + "    {\n"
             + "      \"ml_inference\": {\n"
             + "        \"tag\": \"ml_inference\",\n"
@@ -449,7 +448,7 @@ public void testMLInferenceProcessorRemoteModelOptionalInputs() throws Exception
             + "\",\n"
             + "        \"optional_input_map\": [\n"
             + "          {\n"
-            + "            \"inputText\": \"diary\",\n"
+            + "            \"inputText\": \"diary[0]\",\n"
             + "            \"inputImage\": \"diary_image\"\n"
             + "          }\n"
             + "        ],\n"
@@ -474,7 +473,7 @@ public void testMLInferenceProcessorRemoteModelOptionalInputs() throws Exception
         Map response = searchWithPipeline(client(), index_name, pipelineName, query);
 
         assertEquals((int) JsonPath.parse(response).read("$.hits.hits.length()"), 1);
-        Assert.assertEquals(JsonPath.parse(response).read("$.hits.hits[0]._source.multi_modal_embedding.length()"), "1024");
+        assertEquals((int) JsonPath.parse(response).read("$.hits.hits[0]._source.multi_modal_embedding.length()"), 1024);
     }
 
     /**