Cleanup other resources created as part of Integ Tests

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

Author: dbwiddis

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

@@ -16,12 +16,14 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import org.apache.http.Header;
 import org.apache.http.message.BasicHeader;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.Response;
+import org.opensearch.client.ResponseException;
 import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.core.common.bytes.BytesArray;
 import org.opensearch.core.rest.RestStatus;
@@ -106,11 +108,10 @@ protected static RestRequest getRestRequestWithHeadersAndContent(String tenantId
         if (tenantId != null) {
             headers.put(Constants.TENANT_ID_HEADER, singletonList(tenantId));
         }
-        RestRequest request = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
+        return new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
             .withHeaders(headers)
             .withContent(new BytesArray(requestContent), JSON)
             .build();
-        return request;
     }
 
     @SuppressWarnings("unchecked")
@@ -187,6 +188,39 @@ protected void refreshBeforeSearch(boolean extraDelay) {
         }
     }
 
+    /**
+     * Delete the specified document and wait until a search matches only the specified number of hits
+     * @param tenantId The tenant ID to filter the search by
+     * @param restPath The base path for the REST API
+     * @param id The document ID to be appended to the REST API for deletion
+     * @param hits The number of hits to expect after the deletion is processed
+     * @throws Exception on failures with building or making the request
+     */
+    protected static void deleteAndWaitForSearch(String tenantId, String restPath, String id, int hits) throws Exception {
+        RestRequest request = new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY)
+            .withHeaders(Map.of(TENANT_ID_HEADER, singletonList(tenantId)))
+            .build();
+        // First process the deletion. Dependent resources (e.g. model with connector) may cause 409 status until they are deleted
+        assertBusy(() -> {
+            try {
+                Response deleteResponse = makeRequest(request, DELETE, restPath + id);
+                // first successful deletion should produce an OK
+                assertOK(deleteResponse);
+            } catch (ResponseException e) {
+                // repeat deletions can produce a 404, treat as a success
+                assertNotFound(e.getResponse());
+            }
+        }, 20, TimeUnit.SECONDS);
+        // Deletion processed, now wait for it to disappear from search
+        RestRequest searchRequest = getRestRequestWithHeadersAndContent(tenantId, MATCH_ALL_QUERY);
+        assertBusy(() -> {
+            Response response = makeRequest(searchRequest, GET, restPath + "_search");
+            assertOK(response);
+            SearchResponse searchResponse = searchResponseFromResponse(response);
+            assertEquals(hits, searchResponse.getHits().getTotalHits().value);
+        }, 20, TimeUnit.SECONDS);
+    }
+
     protected static String registerRemoteModelContent(String description, String connectorId, String modelGroupId) {
         StringBuilder sb = new StringBuilder();
         sb.append("{\n");

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

@@ -320,6 +320,12 @@ public void testAgentCRUD() throws Exception {
         assertNotFound(response);
         map = responseToMap(response);
         assertEquals("Failed to find agent with the provided agent id: " + otherAgentId, getErrorReasonFromResponseMap(map));
+
+        /*
+         * Cleanup other resources created
+         */
+        deleteAndWaitForSearch(tenantId, MODELS_PATH, modelId, 0);
+        deleteAndWaitForSearch(tenantId, CONNECTORS_PATH, connectorId, 0);
     }
 
     private static String registerFlowAgentContent(String name, String modelId) {

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

@@ -150,6 +150,7 @@ public void testModelGroupCRUD() throws Exception {
             otherTenantId,
             registerRemoteModelContent("test model", otherConnectorId, sameGroupModelGroupId)
         );
+        String modelInSameGroupOtherTenant = null;
         if (multiTenancyEnabled) {
             ResponseException ex = assertThrows(
                 ResponseException.class,
@@ -168,6 +169,9 @@ public void testModelGroupCRUD() throws Exception {
         } else {
             response = makeRequest(registerModelInSameGroupOtherTenantRequest, POST, MODELS_PATH + "_register");
             assertOK(response);
+            map = responseToMap(response);
+            assertTrue(map.containsKey(MODEL_ID_FIELD));
+            modelInSameGroupOtherTenant = map.get(MODEL_ID_FIELD).toString();
         }
 
         /*
@@ -318,7 +322,7 @@ public void testModelGroupCRUD() throws Exception {
         /*
          * Delete
          */
-        // Delete the models
+        // Delete the model groups
         // First test that we can't delete other tenant model groups
         if (multiTenancyEnabled) {
             ResponseException ex = assertThrows(
@@ -388,6 +392,19 @@ public void testModelGroupCRUD() throws Exception {
             "Failed to find model group with the provided model group id: " + otherModelGroupId,
             getErrorReasonFromResponseMap(map)
         );
+
+        /*
+         * Cleanup other resources created
+         */
+        int additionalModel = modelInSameGroupOtherTenant != null ? 1 : 0;
+        // Need to wait until it's gone from search due to https://github.com/opensearch-project/ml-commons/issues/2932
+        deleteAndWaitForSearch(tenantId, MODELS_PATH, modelId, 1 + additionalModel);
+        deleteAndWaitForSearch(tenantId, MODELS_PATH, sameGroupModelId, additionalModel);
+        if (modelInSameGroupOtherTenant != null) {
+            deleteAndWaitForSearch(otherTenantId, MODELS_PATH, modelInSameGroupOtherTenant, 0);
+        }
+        deleteAndWaitForSearch(tenantId, CONNECTORS_PATH, connectorId, multiTenancyEnabled ? 0 : 1);
+        deleteAndWaitForSearch(otherTenantId, CONNECTORS_PATH, otherConnectorId, 0);
     }
 
     private static String registerModelGroupContent(String name) {

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

@@ -177,6 +177,7 @@ public void testModelCRUD() throws Exception {
             otherTenantId,
             registerRemoteModelContent("other test model", connectorId, null)
         );
+        String wrongTenantModelId = null;
         if (multiTenancyEnabled) {
             ResponseException ex = assertThrows(
                 ResponseException.class,
@@ -194,6 +195,9 @@ public void testModelCRUD() throws Exception {
         } else {
             response = makeRequest(wrongTenantModelRequest, POST, MODELS_PATH + "_register");
             assertOK(response);
+            map = responseToMap(response);
+            assertTrue(map.containsKey(MODEL_ID_FIELD));
+            wrongTenantModelId = map.get(MODEL_ID_FIELD).toString();
         }
 
         // Now register a model with correct connector
@@ -334,5 +338,27 @@ public void testModelCRUD() throws Exception {
         assertNotFound(response);
         map = responseToMap(response);
         assertEquals("Failed to find model with the provided model id: " + otherModelId, getErrorReasonFromResponseMap(map));
+
+        // Delete the model created with null tenant if applicable
+        if (wrongTenantModelId != null) {
+            response = makeRequest(otherTenantRequest, DELETE, MODELS_PATH + wrongTenantModelId);
+            assertOK(response);
+            map = responseToMap(response);
+            assertEquals(wrongTenantModelId, map.get(DOC_ID).toString());
+
+            // Verify the deletion
+            final String deletedModelId = wrongTenantModelId;
+            ex = assertThrows(ResponseException.class, () -> makeRequest(otherTenantRequest, GET, MODELS_PATH + deletedModelId));
+            response = ex.getResponse();
+            assertNotFound(response);
+            map = responseToMap(response);
+            assertEquals("Failed to find model with the provided model id: " + deletedModelId, getErrorReasonFromResponseMap(map));
+        }
+
+        /*
+         * Cleanup other resources created
+         */
+        deleteAndWaitForSearch(tenantId, CONNECTORS_PATH, connectorId, multiTenancyEnabled ? 0 : 1);
+        deleteAndWaitForSearch(otherTenantId, CONNECTORS_PATH, otherConnectorId, 0);
     }
 }