fix stub

temporalio · Oct 9, 2024 · 18bde40 · 18bde40
1 parent 1daf458
commit 18bde40
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java
@@ -266,9 +266,15 @@ public void respondWorkflowTaskFailureMetricTest() throws Exception {
             eagerActivityDispatcher,
             slotSupplier);
 
+    WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub =
+        mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class);
+    when(futureStub.shutdownWorker(any(ShutdownWorkerRequest.class)))
+        .thenReturn(Futures.immediateFuture(ShutdownWorkerResponse.newBuilder().build()));
+
     WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub =
         mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
     when(client.blockingStub()).thenReturn(blockingStub);
+    when(client.futureStub()).thenReturn(futureStub);
     when(blockingStub.withOption(any(), any())).thenReturn(blockingStub);
 
     PollWorkflowTaskQueueResponse pollResponse =