corrected tests

apache · Dec 9, 2024 · 2e78702 · 2e78702
1 parent 4503325
commit 2e78702
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 19 deletions.
diff --git a/gobblin-temporal/src/main/java/org/apache/gobblin/temporal/yarn/YarnService.java b/gobblin-temporal/src/main/java/org/apache/gobblin/temporal/yarn/YarnService.java
@@ -194,7 +194,7 @@ class YarnService extends AbstractIdleService {
   private volatile boolean shutdownInProgress = false;
 
   private final boolean jarCacheEnabled;
-  private final long DEFAULT_ALLOCATION_REQUEST_ID = 0L;
+  private static final long DEFAULT_ALLOCATION_REQUEST_ID = 0L;
   private final AtomicLong allocationRequestIdGenerator = new AtomicLong(DEFAULT_ALLOCATION_REQUEST_ID);
   private final ConcurrentMap<Long, WorkerProfile> workerProfileByAllocationRequestId = new ConcurrentHashMap<>();
 
@@ -555,11 +555,13 @@ protected ByteBuffer getSecurityTokens() throws IOException {
   @VisibleForTesting
   protected String buildContainerCommand(Container container, String helixParticipantId, String helixInstanceTag) {
     long allocationRequestId = container.getAllocationRequestId();
-    // Using getOrDefault for backward-compatibility with containers that don't have allocationRequestId set
     WorkerProfile workerProfile = Optional.fromNullable(this.workerProfileByAllocationRequestId.get(allocationRequestId))
         .or(() -> {
-          LOGGER.warn("No Worker Profile found for {} ... falling back... ", allocationRequestId);
-          return this.workerProfileByAllocationRequestId.get(DEFAULT_ALLOCATION_REQUEST_ID);
+          LOGGER.warn("No Worker Profile found for {}, so falling back to default", allocationRequestId);
+          return this.workerProfileByAllocationRequestId.computeIfAbsent(DEFAULT_ALLOCATION_REQUEST_ID, k -> {
+            LOGGER.warn("WARNING: (LIKELY) UNEXPECTED CONCURRENCY: No Worker Profile even yet mapped to the default allocation request ID {} - creating one now", DEFAULT_ALLOCATION_REQUEST_ID);
+            return new WorkerProfile(this.config);
+          });
         });
     Config workerProfileConfig = workerProfile.getConfig();
 

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
@@ -54,13 +54,11 @@ public void setup() {
   @Test
   public void testReviseWorkforcePlanAndRequestNewContainers() throws Exception {
     int numNewContainers = 5;
-    DynamicScalingYarnService dynamicScalingYarnService = new DynamicScalingYarnService(this.defaultConfigs, "testApp", "testAppId", yarnConfiguration, mockFileSystem, eventBus) {
-      @Override
-      protected void requestContainers(int numContainers, Resource resource, Optional<Long> allocationRequestId) {
-        Assert.assertEquals(numContainers, numNewContainers);
-      }
-    };
+    DynamicScalingYarnService dynamicScalingYarnService = new DynamicScalingYarnService(this.defaultConfigs, "testApp", "testAppId", yarnConfiguration, mockFileSystem, eventBus);
+    DynamicScalingYarnService dynamicScalingYarnServiceSpy = Mockito.spy(dynamicScalingYarnService);
+    Mockito.doNothing().when(dynamicScalingYarnServiceSpy).requestContainers(Mockito.anyInt(), Mockito.any(Resource.class), Mockito.any(Optional.class));
     ScalingDirective baseScalingDirective = new ScalingDirective(WorkforceProfiles.BASELINE_NAME, numNewContainers, System.currentTimeMillis());
-    dynamicScalingYarnService.reviseWorkforcePlanAndRequestNewContainers(Collections.singletonList(baseScalingDirective));
+    dynamicScalingYarnServiceSpy.reviseWorkforcePlanAndRequestNewContainers(Collections.singletonList(baseScalingDirective));
+    Mockito.verify(dynamicScalingYarnServiceSpy, Mockito.times(1)).requestContainers(Mockito.eq(numNewContainers), Mockito.any(Resource.class), Mockito.any(Optional.class));
   }
 }
diff --git a/gobblin-temporal/src/test/java/org/apache/gobblin/temporal/yarn/YarnServiceTest.java b/gobblin-temporal/src/test/java/org/apache/gobblin/temporal/yarn/YarnServiceTest.java
@@ -83,14 +83,13 @@ public void setup() throws IOException, YarnException {
 
   @Test
   public void testYarnServiceStartupWithInitialContainers() throws Exception {
-    int expectedNumContainers = 0;
-    YarnService yarnService = new YarnService(this.defaultConfigs, "testApplicationName", "testApplicationId", yarnConfiguration, mockFileSystem, eventBus) {
-      @Override
-      protected void requestContainers(int numContainers, Resource resource, Optional<Long> allocationRequestId) {
-        Assert.assertEquals(numContainers, expectedNumContainers);
-      }
-    };
-    yarnService.startUp();
+    int expectedNumContainers = 3;
+    Config config = this.defaultConfigs.withValue(GobblinYarnConfigurationKeys.INITIAL_CONTAINERS_KEY, ConfigValueFactory.fromAnyRef(expectedNumContainers));
+    YarnService yarnService = new YarnService(config, "testApplicationName", "testApplicationId", yarnConfiguration, mockFileSystem, eventBus);
+    YarnService yarnServiceSpy = Mockito.spy(yarnService);
+    Mockito.doNothing().when(yarnServiceSpy).requestContainers(Mockito.anyInt(), Mockito.any(Resource.class), Mockito.any(Optional.class));
+    yarnServiceSpy.startUp();
+    Mockito.verify(yarnServiceSpy, Mockito.times(1)).requestContainers(Mockito.eq(expectedNumContainers), Mockito.any(Resource.class), Mockito.any(Optional.class));
   }
 
   @Test