Add get model test case with filter

Added test case to make sure output from get model can be reduced by adding filter_path parameter. Signed-off-by: Vijayan Balasubramanian <[email protected]>
opensearch-project · Oct 28, 2021 · c9a6651 · c9a6651
1 parent 769a137
commit c9a6651
Showing 1 changed file with 41 additions and 0 deletions.
diff --git a/src/test/java/org/opensearch/knn/plugin/action/RestGetModelHandlerIT.java b/src/test/java/org/opensearch/knn/plugin/action/RestGetModelHandlerIT.java
@@ -11,6 +11,7 @@
 
 package org.opensearch.knn.plugin.action;
 
+import joptsimple.internal.Strings;
 import org.apache.http.util.EntityUtils;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
@@ -42,6 +43,8 @@
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
 import java.time.format.DateTimeFormatter;
+import java.util.Arrays;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
@@ -105,6 +108,44 @@ public void testGetModelExists() throws IOException {
         assertEquals(testModelMetadata.getTimestamp(), responseMap.get(MODEL_TIMESTAMP));
     }
 
+
+
+    public void testGetModelExistsWithFilter() throws IOException {
+        createModelSystemIndex();
+        String testModelID = "test-model-id";
+        byte[] testModelBlob = "hello".getBytes();
+        ModelMetadata testModelMetadata = getModelMetadata();
+
+        addModelToSystemIndex(testModelID, testModelMetadata, testModelBlob);
+
+        String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, testModelID);
+        Request request = new Request("GET", restURI);
+
+        List<String> filterdPath = Arrays.asList(MODEL_ID,MODEL_DESCRIPTION,MODEL_TIMESTAMP,KNN_ENGINE);
+        request.addParameter("filter_path", Strings.join(filterdPath, ","));
+
+        Response response = client().performRequest(request);
+        assertEquals(RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
+
+        String responseBody = EntityUtils.toString(response.getEntity());
+        assertNotNull(responseBody);
+
+        Map<String, Object> responseMap = createParser(
+            XContentType.JSON.xContent(),
+            responseBody
+        ).map();
+
+        assertTrue(responseMap.size()==filterdPath.size());
+        assertEquals(testModelID, responseMap.get(MODEL_ID));
+        assertEquals(testModelMetadata.getDescription(), responseMap.get(MODEL_DESCRIPTION));
+        assertEquals(testModelMetadata.getTimestamp(), responseMap.get(MODEL_TIMESTAMP));
+        assertEquals(testModelMetadata.getKnnEngine().getName(), responseMap.get(KNN_ENGINE));
+        assertFalse(responseMap.containsKey(DIMENSION));
+        assertFalse(responseMap.containsKey(MODEL_ERROR));
+        assertFalse(responseMap.containsKey(METHOD_PARAMETER_SPACE_TYPE));
+        assertFalse(responseMap.containsKey(MODEL_STATE));
+    }
+
     public void testGetModelFailsInvalid() throws IOException {
         createModelSystemIndex();
         String restURI = String.join("/", KNNPlugin.KNN_BASE_URI, MODELS, "invalid-model-id");