fix ut

Signed-off-by: Yaliang Wu <ylwu@amazon.com>

Author: ylwu-amzn

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/AwsConnectorExecutorTest.java

@@ -245,7 +245,7 @@ public void executePredict_TextDocsInferenceInput_withStepSize() {
         AwsConnectorExecutor executor = spy(new AwsConnectorExecutor(connector));
         Settings settings = Settings.builder().build();
         threadContext = new ThreadContext(settings);
-        when(executor.getClient()).thenReturn(client);
+        executor.setClient(client);
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
 
@@ -724,7 +724,7 @@ public void executePredict_whenRetryEnabled_thenInvokeRemoteServiceWithRetry() {
         Settings settings = Settings.builder().build();
         threadContext = new ThreadContext(settings);
         ExecutorService executorService = mock(ExecutorService.class);
-        when(executor.getClient()).thenReturn(client);
+        executor.setClient(client);
         when(client.threadPool()).thenReturn(threadPool);
         when(threadPool.getThreadContext()).thenReturn(threadContext);
         when(threadPool.executor(any())).thenReturn(executorService);

ml-algorithms/src/test/java/org/opensearch/ml/engine/algorithms/remote/HttpJsonConnectorExecutorTest.java

@@ -29,6 +29,8 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 import org.opensearch.common.collect.Tuple;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.concurrent.ThreadContext;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.ml.common.FunctionName;
 import org.opensearch.ml.common.connector.Connector;
@@ -41,6 +43,8 @@
 import org.opensearch.ml.common.output.model.ModelTensors;
 import org.opensearch.ml.common.transport.MLTaskResponse;
 import org.opensearch.ml.engine.algorithms.remote.streaming.StreamPredictActionListener;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.client.Client;
 
 import com.google.common.collect.ImmutableMap;
 
@@ -51,9 +55,19 @@ public class HttpJsonConnectorExecutorTest {
     @Mock
     private ActionListener<Tuple<Integer, ModelTensors>> actionListener;
 
+    @Mock
+    private ThreadPool threadPool;
+
+    @Mock
+    private Client client;
+
+    private ThreadContext threadContext;
+
     @Before
     public void setUp() {
         MockitoAnnotations.openMocks(this);
+        Settings settings = Settings.builder().build();
+        threadContext = new ThreadContext(settings);
     }
 
     @Test
@@ -95,8 +109,11 @@ public void invokeRemoteService_invalidIpAddress() {
             .protocol("http")
             .actions(Arrays.asList(predictAction))
             .build();
-        HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
         executor.setConnectorPrivateIpEnabled(false);
+        executor.setClient(client);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -128,8 +145,11 @@ public void invokeRemoteService_EnabledPrivateIpAddress() {
             .protocol("http")
             .actions(Arrays.asList(predictAction))
             .build();
-        HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
         executor.setConnectorPrivateIpEnabled(true);
+        executor.setClient(client);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -158,8 +178,11 @@ public void invokeRemoteService_DisabledPrivateIpAddress() {
             .protocol("http")
             .actions(Arrays.asList(predictAction))
             .build();
-        HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
         executor.setConnectorPrivateIpEnabled(false);
+        executor.setClient(client);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         executor
             .invokeRemoteService(
                 PREDICT.name(),
@@ -215,7 +238,10 @@ public void invokeRemoteService_get_request() {
             .protocol("http")
             .actions(Arrays.asList(predictAction))
             .build();
-        HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
+        executor.setClient(client);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         executor.invokeRemoteService(PREDICT.name(), createMLInput(), new HashMap<>(), null, new ExecutionContext(0), actionListener);
     }
 
@@ -235,7 +261,10 @@ public void invokeRemoteService_post_request() {
             .protocol("http")
             .actions(Arrays.asList(predictAction))
             .build();
-        HttpJsonConnectorExecutor executor = new HttpJsonConnectorExecutor(connector);
+        HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
+        executor.setClient(client);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         executor
             .invokeRemoteService(PREDICT.name(), createMLInput(), new HashMap<>(), "hello world", new ExecutionContext(0), actionListener);
     }
@@ -257,6 +286,9 @@ public void invokeRemoteService_nullHttpClient_throwMLException() {
             .actions(Arrays.asList(predictAction))
             .build();
         HttpJsonConnectorExecutor executor = spy(new HttpJsonConnectorExecutor(connector));
+        executor.setClient(client);
+        when(client.threadPool()).thenReturn(threadPool);
+        when(threadPool.getThreadContext()).thenReturn(threadContext);
         when(executor.getHttpClient()).thenReturn(null);
         executor
             .invokeRemoteService(PREDICT.name(), createMLInput(), new HashMap<>(), "hello world", new ExecutionContext(0), actionListener);