fix ut

Signed-off-by: Yaliang Wu <[email protected]>
ylwu-amzn · Sep 14, 2024 · 820b438 · 820b438
1 parent 784bcb8
commit 820b438
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/plugin/src/test/java/org/opensearch/ml/action/tasks/GetTaskTransportActionTests.java b/plugin/src/test/java/org/opensearch/ml/action/tasks/GetTaskTransportActionTests.java
@@ -15,12 +15,18 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_CANCELLED_REGEX;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_CANCELLING_REGEX;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_COMPLETED_REGEX;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_EXPIRED_REGEX;
+import static org.opensearch.ml.settings.MLCommonsSettings.ML_COMMONS_REMOTE_JOB_STATUS_FIELD;
 
 import java.io.IOException;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.junit.Before;
 import org.junit.Ignore;
@@ -36,6 +42,7 @@
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.Metadata;
 import org.opensearch.cluster.service.ClusterService;
+import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.xcontent.XContentFactory;
@@ -118,7 +125,14 @@ public void setup() throws IOException {
         MockitoAnnotations.openMocks(this);
         mlTaskGetRequest = MLTaskGetRequest.builder().taskId("test_id").build();
 
-        Settings settings = Settings.builder().build();
+        Settings settings = Settings
+            .builder()
+            .putList(ML_COMMONS_REMOTE_JOB_STATUS_FIELD.getKey(), List.of("status", "TransformJobStatus"))
+            .put(ML_COMMONS_REMOTE_JOB_STATUS_COMPLETED_REGEX.getKey(), "(complete|completed)")
+            .put(ML_COMMONS_REMOTE_JOB_STATUS_CANCELLED_REGEX.getKey(), "(stopped|cancelled)")
+            .put(ML_COMMONS_REMOTE_JOB_STATUS_CANCELLING_REGEX.getKey(), "(stopping|cancelling)")
+            .put(ML_COMMONS_REMOTE_JOB_STATUS_EXPIRED_REGEX.getKey(), "(expired|timeout)")
+            .build();
         threadContext = new ThreadContext(settings);
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
@@ -127,6 +141,21 @@ public void setup() throws IOException {
         doReturn(metaData).when(clusterState).metadata();
 
         doReturn(true).when(metaData).hasIndex(anyString());
+        when(clusterService.getSettings()).thenReturn(settings);
+        when(this.clusterService.getClusterSettings())
+            .thenReturn(
+                new ClusterSettings(
+                    settings,
+                    Set
+                        .of(
+                            ML_COMMONS_REMOTE_JOB_STATUS_FIELD,
+                            ML_COMMONS_REMOTE_JOB_STATUS_COMPLETED_REGEX,
+                            ML_COMMONS_REMOTE_JOB_STATUS_CANCELLED_REGEX,
+                            ML_COMMONS_REMOTE_JOB_STATUS_CANCELLING_REGEX,
+                            ML_COMMONS_REMOTE_JOB_STATUS_EXPIRED_REGEX
+                        )
+                )
+            );
 
         getTaskTransportAction = spy(
             new GetTaskTransportAction(

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLBatchIngestionActionTests.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLBatchIngestionActionTests.java
@@ -11,7 +11,6 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
-import static org.opensearch.ml.rest.RestMLCancelBatchJobAction.ML_CANCEL_TASK_ACTION;
 import static org.opensearch.ml.utils.TestHelper.getBatchIngestionRestRequest;
 
 import java.io.IOException;

diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLCancelBatchJobActionTests.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLCancelBatchJobActionTests.java
@@ -74,7 +74,7 @@ public void testConstructor() {
     public void testGetName() {
         String actionName = restMLCancelBatchJobAction.getName();
         assertFalse(Strings.isNullOrEmpty(actionName));
-        assertEquals("ml_cancel_batch_action", actionName);
+        assertEquals("ml_cancel_task_action", actionName);
     }
 
     public void testRoutes() {
@@ -83,7 +83,7 @@ public void testRoutes() {
         assertFalse(routes.isEmpty());
         RestHandler.Route route = routes.get(0);
         assertEquals(RestRequest.Method.POST, route.getMethod());
-        assertEquals("/_plugins/_ml/tasks/{task_id}/_cancel_batch", route.getPath());
+        assertEquals("/_plugins/_ml/tasks/{task_id}/_cancel", route.getPath());
     }
 
     public void test_PrepareRequest() throws Exception {