fix UT

Signed-off-by: shuangli-z <[email protected]>

Author: shuangli-z

common/src/test/java/org/opensearch/ml/common/indexInsight/FieldDescriptionTaskTests.java

@@ -142,6 +142,7 @@ public void testRunTask_MLConfigRetrievalFailure() {
         String statisticalContent = "{\"mapping\": {\"field1\": {\"type\": \"text\"}}}";
 
         mockGetSuccess(sdkClient, statisticalContent);
+        mockUpdateSuccess(sdkClient);
         mockMLConfigFailure(client, "Config not found");
 
         task.runTask("tenant-id", listener);

common/src/test/java/org/opensearch/ml/common/indexInsight/LogRelatedIndexCheckTaskTests.java

@@ -277,6 +277,7 @@ public void testRunTask_MLConfigRetrievalFailure() {
         mockSearchResponse();
         // Mock getAgentIdToRun
         mockMLConfigFailure(client, "Config not found");
+        mockUpdateSuccess(sdkClient);
 
         ActionListener<IndexInsight> listener = mock(ActionListener.class);
         task.runTask("tenant-id", listener);

common/src/test/java/org/opensearch/ml/common/indexInsight/StatisticalDataTaskTests.java

@@ -38,6 +38,8 @@
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.cluster.metadata.MappingMetadata;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.common.bytes.BytesReference;
@@ -52,6 +54,7 @@
 import org.opensearch.search.aggregations.bucket.sampler.InternalSampler;
 import org.opensearch.search.aggregations.metrics.InternalTopHits;
 import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.AdminClient;
 import org.opensearch.transport.client.Client;
 import org.opensearch.transport.client.IndicesAdminClient;
@@ -68,9 +71,14 @@ private Client setupBasicClientMocks() {
         Client client = mock(Client.class);
         AdminClient adminClient = mock(AdminClient.class);
         IndicesAdminClient indicesAdminClient = mock(IndicesAdminClient.class);
+        ThreadPool threadPool = mock(ThreadPool.class);
+        Settings settings = Settings.builder().build();
+        ThreadContext threadContext = new ThreadContext(settings);
 
         when(client.admin()).thenReturn(adminClient);
         when(adminClient.indices()).thenReturn(indicesAdminClient);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
 
         return client;
     }
@@ -184,6 +192,8 @@ public void testRunTask_WithEmptyMappings() {
 
         when(getMappingsResponse.getMappings()).thenReturn(new HashMap<>());
         setupGetMappingsCall(client, getMappingsResponse);
+        sdkClient = mock(SdkClient.class);
+        mockUpdateSuccess(sdkClient);
 
         StatisticalDataTask task = new StatisticalDataTask("test-index", client, sdkClient);
         task.runTask("tenant-id", listener);