adding immediate refresh to delete model group request (opensearch-pr…

…oject#2514) * adding immediate refresh to delete model group request Signed-off-by: Bhavana Ramaram <[email protected]> * fix format violations Signed-off-by: Bhavana Ramaram <[email protected]> * add IT tests Signed-off-by: Bhavana Ramaram <[email protected]> * remove thread sleep Signed-off-by: Bhavana Ramaram <[email protected]> --------- Signed-off-by: Bhavana Ramaram <[email protected]>
rbhavna · Jun 10, 2024 · 2b98d20 · 2b98d20
1 parent a0272f2
commit 2b98d20
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 16 deletions.
diff --git a/...n/src/main/java/org/opensearch/ml/action/model_group/DeleteModelGroupTransportAction.java b/...n/src/main/java/org/opensearch/ml/action/model_group/DeleteModelGroupTransportAction.java
@@ -15,6 +15,7 @@
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.support.ActionFilters;
 import org.opensearch.action.support.HandledTransportAction;
+import org.opensearch.action.support.WriteRequest;
 import org.opensearch.client.Client;
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.inject.Inject;
@@ -68,7 +69,8 @@ public DeleteModelGroupTransportAction(
     protected void doExecute(Task task, ActionRequest request, ActionListener<DeleteResponse> actionListener) {
         MLModelGroupDeleteRequest mlModelGroupDeleteRequest = MLModelGroupDeleteRequest.fromActionRequest(request);
         String modelGroupId = mlModelGroupDeleteRequest.getModelGroupId();
-        DeleteRequest deleteRequest = new DeleteRequest(ML_MODEL_GROUP_INDEX, modelGroupId);
+        DeleteRequest deleteRequest = new DeleteRequest(ML_MODEL_GROUP_INDEX, modelGroupId)
+            .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
         User user = RestActionUtils.getUserContext(client);
         try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
             ActionListener<DeleteResponse> wrappedListener = ActionListener.runBefore(actionListener, () -> context.restore());

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLDeleteModelGroupActionIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLDeleteModelGroupActionIT.java
@@ -6,32 +6,53 @@
 package org.opensearch.ml.rest;
 
 import java.io.IOException;
-import java.util.Map;
 
-import org.apache.hc.core5.http.HttpEntity;
-import org.junit.Ignore;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.rules.ExpectedException;
 import org.opensearch.client.Response;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.ml.common.transport.model_group.MLRegisterModelGroupInput;
 import org.opensearch.ml.utils.TestHelper;
 
 public class RestMLDeleteModelGroupActionIT extends MLCommonsRestTestCase {
     @Rule
     public ExpectedException exceptionRule = ExpectedException.none();
+    private MLRegisterModelGroupInput mlRegisterModelGroupInput;
+    private String modelGroupId;
+
+    @Before
+    public void setup() throws IOException {
+        mlRegisterModelGroupInput = MLRegisterModelGroupInput.builder().name("testGroupID").description("This is test Group").build();
+        registerModelGroup(client(), TestHelper.toJsonString(mlRegisterModelGroupInput), registerModelGroupResult -> {
+            this.modelGroupId = (String) registerModelGroupResult.get("model_group_id");
+        });
+    }
 
-    @Ignore
     public void testDeleteModelGroupAPI_Success() throws IOException {
-        Response trainModelGroupResponse = ingestModelData();
-        HttpEntity entity = trainModelGroupResponse.getEntity();
-        assertNotNull(trainModelGroupResponse);
-        String entityString = TestHelper.httpEntityToString(entity);
-        Map map = gson.fromJson(entityString, Map.class);
-        String model_group_id = (String) map.get("model_group_id");
-
-        Response deleteModelResponse = TestHelper
-            .makeRequest(client(), "DELETE", "/_plugins/_ml/model_groups/" + model_group_id, null, "", null);
-        assertNotNull(deleteModelResponse);
-        assertEquals(RestStatus.OK, TestHelper.restStatus(deleteModelResponse));
+
+        Response deleteModelGroupResponse = TestHelper
+            .makeRequest(client(), "DELETE", "/_plugins/_ml/model_groups/" + modelGroupId, null, "", null);
+        assertNotNull(deleteModelGroupResponse);
+        assertEquals(RestStatus.OK, TestHelper.restStatus(deleteModelGroupResponse));
+    }
+
+    public void testDeleteAndRegisterModelGroup_Success() throws IOException {
+
+        Response deleteModelGroupResponse = TestHelper
+            .makeRequest(client(), "DELETE", "/_plugins/_ml/model_groups/" + modelGroupId, null, "", null);
+
+        if (TestHelper.restStatus(deleteModelGroupResponse).equals(RestStatus.OK)) {
+            MLRegisterModelGroupInput newMlRegisterModelGroupInput = MLRegisterModelGroupInput
+                .builder()
+                .name("testGroupID")
+                .description("This is a new test Group")
+                .build();
+
+            registerModelGroup(client(), TestHelper.toJsonString(newMlRegisterModelGroupInput), registerModelGroupResponse -> {
+                assertNotNull(registerModelGroupResponse);
+                assertEquals("CREATED", registerModelGroupResponse.get("status"));
+            });
+        }
     }
 }