Update mock code for exception test

apache · Sep 5, 2024 · 2162167 · 2162167
1 parent f3f61d6
commit 2162167
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/...src/test/java/org/apache/gobblin/temporal/joblauncher/GobblinTemporalJobLauncherTest.java b/...src/test/java/org/apache/gobblin/temporal/joblauncher/GobblinTemporalJobLauncherTest.java
@@ -24,6 +24,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 
 import org.apache.hadoop.fs.Path;
+import org.mockito.Mock;
 import org.mockito.MockedStatic;
 import org.mockito.Mockito;
 import org.testng.annotations.BeforeClass;
@@ -82,7 +83,6 @@ protected void submitJob(List<WorkUnit> workUnits)
   public void setUp() throws Exception {
     mockServiceStubs = mock(WorkflowServiceStubs.class);
     mockClient = mock(WorkflowClient.class);
-    mockStub = mock(WorkflowStub.class);
     mockExecutionInfo = mock(WorkflowExecutionInfo.class);
     DescribeWorkflowExecutionResponse mockResponse = mock(DescribeWorkflowExecutionResponse.class);
     WorkflowServiceGrpc.WorkflowServiceBlockingStub mockBlockingStub = mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class);
@@ -95,8 +95,6 @@ public void setUp() throws Exception {
         .thenReturn(mockServiceStubs);
     mockWorkflowClientFactory.when(() -> TemporalWorkflowClientFactory.createClientInstance(Mockito.any(), Mockito.anyString()))
         .thenReturn(mockClient);
-    when(mockClient.newUntypedWorkflowStub(Mockito.anyString())).thenReturn(mockStub);
-    when(mockStub.getExecution()).thenReturn(WorkflowExecution.getDefaultInstance());
 
     jobProperties = new Properties();
     jobProperties.setProperty(ConfigurationKeys.FS_URI_KEY, "file:///");
@@ -107,6 +105,10 @@ public void setUp() throws Exception {
 
   @BeforeMethod
   public void methodSetUp() throws Exception {
+    mockStub = mock(WorkflowStub.class);
+    when(mockClient.newUntypedWorkflowStub(Mockito.anyString())).thenReturn(mockStub);
+    when(mockStub.getExecution()).thenReturn(WorkflowExecution.getDefaultInstance());
+
     File tmpDir = Files.createTempDir();
     String basePath = tmpDir.getAbsolutePath();
     Path appWorkDir = new Path(basePath, "testAppWorkDir");
@@ -178,5 +180,7 @@ public void testCancelWorkflowFetchStatusThrowsException() throws Exception {
     jobLauncher.executeCancellation();
 
     verify(mockStub, times(1)).cancel();
+
+    Mockito.reset(mockExecutionInfo);
   }
 }