diff --git a/src/test/java/org/opensearch/knn/plugin/action/RestGetModelHandlerIT.java b/src/test/java/org/opensearch/knn/plugin/action/RestGetModelHandlerIT.java index d982cd3ce..c45452611 100644 --- 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 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 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");