add ITs

Signed-off-by: Mingshi Liu <[email protected]>
opensearch-project · Jul 22, 2024 · f8edd8c · f8edd8c
1 parent bf4eda2
commit f8edd8c
Show file tree

Hide file tree

Showing 7 changed files with 1,603 additions and 177 deletions.
diff --git a/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchResponseProcessor.java b/plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchResponseProcessor.java
diff --git a/...pensearch/ml/utils/VersionedMapUtils.java → ...ava/org/opensearch/ml/utils/MapUtils.java b/...pensearch/ml/utils/VersionedMapUtils.java → ...ava/org/opensearch/ml/utils/MapUtils.java
@@ -8,7 +8,7 @@
 import java.util.HashMap;
 import java.util.Map;
 
-public class VersionedMapUtils {
+public class MapUtils {
 
     /**
      * Increments the counter for the given key in the specified version.
@@ -35,4 +35,19 @@ public static int getCounter(Map<Integer, Map<String, Integer>> versionedCounter
         return counters != null ? counters.getOrDefault(key, -1) : 0;
     }
 
+    /**
+     * Increments the counter value for the given key in the provided counters map.
+     * If the key does not exist in the map, it is added with an initial counter value of 0.
+     *
+     * @param counters A map that stores integer counters for each integer key.
+     * @param key The integer key for which the counter needs to be incremented.
+     */
+    public static void incrementCounter(Map<Integer, Integer> counters, int key) {
+        counters.put(key, counters.getOrDefault(key, 0) + 1);
+    }
+
+    public static int getCounter(Map<Integer, Integer> counters, int key) {
+        return counters.getOrDefault(key, 0);
+    }
+
 }
diff --git a/plugin/src/test/java/org/opensearch/ml/plugin/MachineLearningPluginTests.java b/plugin/src/test/java/org/opensearch/ml/plugin/MachineLearningPluginTests.java
@@ -38,6 +38,7 @@
 import org.opensearch.ml.common.spi.MLCommonsExtension;
 import org.opensearch.ml.common.spi.tools.Tool;
 import org.opensearch.ml.engine.tools.MLModelTool;
+import org.opensearch.ml.processor.MLInferenceSearchResponseProcessor;
 import org.opensearch.plugins.ExtensiblePlugin;
 import org.opensearch.plugins.SearchPipelinePlugin;
 import org.opensearch.plugins.SearchPlugin;
@@ -83,10 +84,11 @@ public void testGetRequestProcessors() {
     public void testGetResponseProcessors() {
         SearchPipelinePlugin.Parameters parameters = mock(SearchPipelinePlugin.Parameters.class);
         Map<String, ?> responseProcessors = plugin.getResponseProcessors(parameters);
-        assertEquals(1, responseProcessors.size());
+        assertEquals(2, responseProcessors.size());
         assertTrue(
             responseProcessors.get(GenerativeQAProcessorConstants.RESPONSE_PROCESSOR_TYPE) instanceof GenerativeQAResponseProcessor.Factory
         );
+        assertTrue(responseProcessors.get(MLInferenceSearchResponseProcessor.TYPE) instanceof MLInferenceSearchResponseProcessor.Factory);
     }
 
     @Test