1818package org .apache .gobblin .temporal .yarn ;
1919
2020import java .util .Collections ;
21-
2221import java .util .List ;
22+
2323import org .apache .hadoop .fs .FileSystem ;
2424import org .apache .hadoop .yarn .api .protocolrecords .RegisterApplicationMasterResponse ;
2525import org .apache .hadoop .yarn .api .records .ApplicationAttemptId ;
@@ -170,17 +170,12 @@ public void testHandleContainerCompletionForStatusOOM(int containerExitStatusCod
170170 ContainerStatus containerStatus = Mockito .mock (ContainerStatus .class );
171171 Mockito .when (containerStatus .getContainerId ()).thenReturn (containerId );
172172 Mockito .when (containerStatus .getExitStatus ()).thenReturn (containerExitStatusCode );
173-
174- dynamicScalingYarnServiceSpy .startUp ();
175173 dynamicScalingYarnServiceSpy .containerMap .put (containerId , containerInfo ); // Required to be done for test otherwise containerMap is always empty since it is updated after containers are allocated
176-
177174 dynamicScalingYarnServiceSpy .handleContainerCompletion (containerStatus );
178-
179175 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (1 )).reviseWorkforcePlanAndRequestNewContainers (Mockito .anyList ());
180176 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (2 )).requestContainersForWorkerProfile (Mockito .any (WorkerProfile .class ), Mockito .anyInt ());
181177 ArgumentCaptor <Resource > resourceCaptor = ArgumentCaptor .forClass (Resource .class );
182178 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (2 )).requestContainers (Mockito .eq (1 ), resourceCaptor .capture (), Mockito .any (Optional .class ));
183-
184179 Resource capturedResource = resourceCaptor .getValue ();
185180 Assert .assertEquals (capturedResource .getMemorySize (), (long ) initMemoryMbs * DynamicScalingYarnService .DEFAULT_REPLACEMENT_CONTAINER_MEMORY_MULTIPLIER );
186181 Assert .assertEquals (capturedResource .getVirtualCores (), initCores );
@@ -193,16 +188,12 @@ public void testHandleContainerCompletionForNonOOMStatusWhichRequestReplacementC
193188 ContainerStatus containerStatus = Mockito .mock (ContainerStatus .class );
194189 Mockito .when (containerStatus .getContainerId ()).thenReturn (containerId );
195190 Mockito .when (containerStatus .getExitStatus ()).thenReturn (containerExitStatusCode );
196-
197- dynamicScalingYarnServiceSpy .startUp ();
198191 dynamicScalingYarnServiceSpy .containerMap .put (containerId , containerInfo ); // Required to be done for test otherwise containerMap is always empty since it is updated after containers are allocated
199-
200192 dynamicScalingYarnServiceSpy .handleContainerCompletion (containerStatus );
201193 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (0 )).reviseWorkforcePlanAndRequestNewContainers (Mockito .anyList ());
202- Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (2 )).requestContainersForWorkerProfile (Mockito .any (WorkerProfile .class ), Mockito .anyInt ());
194+ Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (1 )).requestContainersForWorkerProfile (Mockito .any (WorkerProfile .class ), Mockito .anyInt ());
203195 ArgumentCaptor <Resource > resourceCaptor = ArgumentCaptor .forClass (Resource .class );
204- Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (2 )).requestContainers (Mockito .eq (1 ), resourceCaptor .capture (), Mockito .any (Optional .class ));
205-
196+ Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (1 )).requestContainers (Mockito .eq (1 ), resourceCaptor .capture (), Mockito .any (Optional .class ));
206197 Resource capturedResource = resourceCaptor .getValue ();
207198 Assert .assertEquals (capturedResource .getMemorySize (), initMemoryMbs );
208199 Assert .assertEquals (capturedResource .getVirtualCores (), initCores );
@@ -230,7 +221,6 @@ public void testHandleContainerCompletionForAllOOMStatus() throws Exception {
230221 Mockito .when (containerStatus3 .getContainerId ()).thenReturn (containerId3 );
231222 Mockito .when (containerStatus3 .getExitStatus ()).thenReturn (ContainerExitStatus .KILLED_EXCEEDED_PMEM );
232223
233- dynamicScalingYarnServiceSpy .startUp ();
234224 // Required to be done for test otherwise containerMap is always empty since it is updated after containers are allocated
235225 dynamicScalingYarnServiceSpy .containerMap .put (containerId1 , containerInfo1 );
236226 dynamicScalingYarnServiceSpy .containerMap .put (containerId2 , containerInfo2 );
@@ -268,7 +258,6 @@ public void testHandleContainerCompletionForExitStatusWhichDoesNotRequestReplace
268258 Mockito .when (containerStatus .getExitStatus ()).thenReturn (containerExitStatusCode );
269259 dynamicScalingYarnServiceSpy .containerMap .put (containerId , containerInfo ); // Required to be done for test otherwise containerMap is always empty since it is updated after containers are allocated
270260 dynamicScalingYarnServiceSpy .handleContainerCompletion (containerStatus );
271- // All zero invocation since startup is not called and no new containers should be requested
272261 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (0 )).reviseWorkforcePlanAndRequestNewContainers (Mockito .anyList ());
273262 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (0 )).requestContainersForWorkerProfile (Mockito .any (WorkerProfile .class ), Mockito .anyInt ());
274263 Mockito .verify (dynamicScalingYarnServiceSpy , Mockito .times (0 )).requestContainers (Mockito .anyInt (), Mockito .any (Resource .class ), Mockito .any (Optional .class ));
0 commit comments