Add an integ test for aggregations in search

Signed-off-by: Daniel Widdis <[email protected]>

Author: dbwiddis

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

@@ -57,6 +57,18 @@ public abstract class MLCommonsTenantAwareRestTestCase extends MLCommonsRestTest
 
     // REST body
     protected static final String MATCH_ALL_QUERY = "{\"query\":{\"match_all\":{}}}";
+    protected static final String AGGREGATION_QUERY = "{\n"
+        + "  \"size\": 0,\n"
+        + "  \"aggs\": {\n"
+        + "    \"unique_model_names\": {\n"
+        + "      \"terms\": {\n"
+        + "        \"field\": \"name.keyword\",\n"
+        + "        \"size\": 10000\n"
+        + "      }\n"
+        + "    }\n"
+        + "  }\n"
+        + "}";
+
     protected static final String EMPTY_CONTENT = "{}";
 
     // REST Response error reasons
@@ -84,6 +96,7 @@ public abstract class MLCommonsTenantAwareRestTestCase extends MLCommonsRestTest
     protected final RestRequest tenantMatchAllRequest = getRestRequestWithHeadersAndContent(tenantId, MATCH_ALL_QUERY);
     protected final RestRequest otherTenantMatchAllRequest = getRestRequestWithHeadersAndContent(otherTenantId, MATCH_ALL_QUERY);
     protected final RestRequest nullTenantMatchAllRequest = getRestRequestWithHeadersAndContent(null, MATCH_ALL_QUERY);
+    protected final RestRequest tenantAggregationRequest = getRestRequestWithHeadersAndContent(tenantId, AGGREGATION_QUERY);
 
     protected static boolean isMultiTenancyEnabled() throws IOException {
         // pass -Dtests.rest.tenantaware=true on gradle command line to enable

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

@@ -16,6 +16,7 @@
 import org.opensearch.client.Response;
 import org.opensearch.client.ResponseException;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.ml.utils.TestHelper;
 import org.opensearch.rest.RestRequest;
 
 public class RestMLModelTenantAwareIT extends MLCommonsTenantAwareRestTestCase {
@@ -277,6 +278,25 @@ public void testModelCRUD() throws Exception {
             assertNull(searchResponse.getHits().getHits()[1].getSourceAsMap().get(TENANT_ID_FIELD));
         }
 
+        // Test search with aggregations to verify no error is thrown
+        Response restResponse = makeRequest(tenantAggregationRequest, GET, MODELS_PATH + "_search");
+        assertOK(restResponse);
+        // We can't (easily) parse the JSON from this response into a SearchResponse for the same reason as the bug we're testing
+        // So we test the string values instead
+        String responseJson = TestHelper.httpEntityToString(restResponse.getEntity());
+        assertTrue(responseJson.contains("\"aggregations\":{\"unique_model_names\""));
+        if (multiTenancyEnabled) {
+            assertTrue(responseJson.contains("\"value\":1,\"relation\":\"eq\""));
+            assertTrue(responseJson.contains("\"key\":\"remote model for connector_id "));
+            assertTrue(responseJson.contains("\"doc_count\":1"));
+        } else {
+            assertTrue(responseJson.contains("\"value\":3,\"relation\":\"eq\""));
+            // one connector will have two doc_count, other will have one
+            assertTrue(responseJson.contains("\"key\":\"remote model for connector_id "));
+            assertTrue(responseJson.contains("\"doc_count\":2"));
+            assertTrue(responseJson.contains("\"doc_count\":1"));
+        }
+
         /*
          * Delete
          */