added test for non-replacement case

apache · Feb 10, 2025 · 729bebb · 729bebb
1 parent 29891fd
commit 729bebb
Showing 1 changed file with 29 additions and 0 deletions.
diff --git a/...emporal/src/test/java/org/apache/gobblin/temporal/yarn/DynamicScalingYarnServiceTest.java b/...emporal/src/test/java/org/apache/gobblin/temporal/yarn/DynamicScalingYarnServiceTest.java
@@ -150,6 +150,19 @@ public Object[][] NonOOMExitStatusProviderWhichRequestReplacementContainer() {
     };
   }
 
+  @DataProvider(name = "ExitStatusProviderWhichDoesNotRequestReplacementContainer")
+  public Object[][] ExitStatusProviderWhichDoesNotRequestReplacementContainer() {
+    return new Object[][] {
+        {ContainerExitStatus.SUCCESS},
+        {ContainerExitStatus.INVALID},
+        {ContainerExitStatus.DISKS_FAILED},
+        {ContainerExitStatus.KILLED_BY_APPMASTER},
+        {ContainerExitStatus.KILLED_BY_RESOURCEMANAGER},
+        {ContainerExitStatus.KILLED_AFTER_APP_COMPLETION},
+        {ContainerExitStatus.KILLED_BY_CONTAINER_SCHEDULER}
+    };
+  }
+
   @Test(dataProvider = "OOMExitStatusProvider")
   public void testHandleContainerCompletionForStatusOOM(int containerExitStatusCode) throws Exception {
     ContainerId containerId = generateRandomContainerId();
@@ -246,6 +259,22 @@ public void testHandleContainerCompletionForAllOOMStatus() throws Exception {
     }
   }
 
+  @Test(dataProvider = "ExitStatusProviderWhichDoesNotRequestReplacementContainer")
+  public void testHandleContainerCompletionForExitStatusWhichDoesNotRequestReplacementContainer(int containerExitStatusCode) throws Exception {
+    ContainerId containerId = generateRandomContainerId();
+    DynamicScalingYarnService.ContainerInfo containerInfo = createBaselineContainerInfo(containerId);
+    ContainerStatus containerStatus = Mockito.mock(ContainerStatus.class);
+    Mockito.when(containerStatus.getContainerId()).thenReturn(containerId);
+    Mockito.when(containerStatus.getExitStatus()).thenReturn(containerExitStatusCode);
+    dynamicScalingYarnServiceSpy.containerMap.put(containerId, containerInfo); // Required to be done for test otherwise containerMap is always empty since it is updated after containers are allocated
+    dynamicScalingYarnServiceSpy.handleContainerCompletion(containerStatus);
+    // All zero invocation since startup is not called and no new containers should be requested
+    Mockito.verify(dynamicScalingYarnServiceSpy, Mockito.times(0)).reviseWorkforcePlanAndRequestNewContainers(Mockito.anyList());
+    Mockito.verify(dynamicScalingYarnServiceSpy, Mockito.times(0)).requestContainersForWorkerProfile(Mockito.any(WorkerProfile.class), Mockito.anyInt());
+    Mockito.verify(dynamicScalingYarnServiceSpy, Mockito.times(0)).requestContainers(Mockito.anyInt(), Mockito.any(Resource.class), Mockito.any(Optional.class));
+  }
+
+
   private ContainerId generateRandomContainerId() {
     return ContainerId.newContainerId(ApplicationAttemptId.newInstance(ApplicationId.newInstance(1, 0),
         0), (long) (Math.random() * 1000));