Repository: apache/flink
File: flink-runtime/src/test/java/org/apache/flink/runtime/resourcemanager/slotmanager/FineGrainedSlotManagerTest.java (Java)

Code before:
                                                            executorConnection1,
                                                            new SlotReport(),
                                                            DEFAULT_TOTAL_RESOURCE_PROFILE,
                                                            DEFAULT_SLOT_RESOURCE_PROFILE)));
                    assertThat(assertFutureCompleteAndReturn(registerTaskManagerFuture1))
                            .isEqualTo(SlotManager.RegistrationResult.SUCCESS);

                    final CompletableFuture<SlotManager.RegistrationResult>
                            registerTaskManagerFuture2 = new CompletableFuture<>();
                    context.runInMainThreadAndWait(
                            () ->
                                    registerTaskManagerFuture2.complete(
                                            context.getSlotManager()
                                                    .registerTaskManager(
                                                            executorConnection2,
                                                            new SlotReport(
                                                                    createAllocatedSlotStatus(
                                                                            new JobID(),
                                                                            new AllocationID(),
                                                                            DEFAULT_SLOT_RESOURCE_PROFILE)),
                                                            DEFAULT_TOTAL_RESOURCE_PROFILE,
                                                            DEFAULT_SLOT_RESOURCE_PROFILE)));
                    assertThat(assertFutureCompleteAndReturn(registerTaskManagerFuture2))
                            .isEqualTo(SlotManager.RegistrationResult.SUCCESS);

                    // triggers the metric update task on the main thread
                    scheduledExecutor.triggerPeriodicScheduledTasks();
                    context.runInMainThreadAndWait(() -> {});

                    assertThat(slotsTotalGauge.get().getValue())
                            .isEqualTo(2 * DEFAULT_NUM_SLOTS_PER_WORKER);
                    assertThat(slotsAvailableGauge.get().getValue())
                            .isEqualTo(2 * DEFAULT_NUM_SLOTS_PER_WORKER - 1);

                    final CompletableFuture<Boolean> unRegisterTaskManagerFuture =
                            new CompletableFuture<>();
                    context.runInMainThreadAndWait(
                            () ->
                                    unRegisterTaskManagerFuture.complete(
                                            context.getSlotManager()
                                                    .unregisterTaskManager(
                                                            executorConnection2.getInstanceID(),
                                                            TEST_EXCEPTION)));
                    assertThat(assertFutureCompleteAndReturn(unRegisterTaskManagerFuture)).isTrue();

                    // triggers the metric update task on the main thread
                    scheduledExecutor.triggerPeriodicScheduledTasks();
                    context.runInMainThreadAndWait(() -> {});

                    assertThat(slotsTotalGauge.get().getValue())
                            .isEqualTo(DEFAULT_NUM_SLOTS_PER_WORKER);

Review comment (suggestion):
```suggestion
                    // triggers the metric update task on the main thread and waits for the main thread to process queued up callbacks
```

Code after:
                                                            executorConnection1,
                                                            new SlotReport(),
                                                            DEFAULT_TOTAL_RESOURCE_PROFILE,
                                                            DEFAULT_SLOT_RESOURCE_PROFILE)));
                    assertThat(assertFutureCompleteAndReturn(registerTaskManagerFuture1))
                            .isEqualTo(SlotManager.RegistrationResult.SUCCESS);

                    final CompletableFuture<SlotManager.RegistrationResult>
                            registerTaskManagerFuture2 = new CompletableFuture<>();
                    context.runInMainThreadAndWait(
                            () ->
                                    registerTaskManagerFuture2.complete(
                                            context.getSlotManager()
                                                    .registerTaskManager(
                                                            executorConnection2,
                                                            new SlotReport(
                                                                    createAllocatedSlotStatus(
                                                                            new JobID(),
                                                                            new AllocationID(),
                                                                            DEFAULT_SLOT_RESOURCE_PROFILE)),
                                                            DEFAULT_TOTAL_RESOURCE_PROFILE,
                                                            DEFAULT_SLOT_RESOURCE_PROFILE)));
                    assertThat(assertFutureCompleteAndReturn(registerTaskManagerFuture2))
                            .isEqualTo(SlotManager.RegistrationResult.SUCCESS);

                    // triggers the metric update task on the main thread and waits for the main
                    // thread to process queued up callbacks
                    scheduledExecutor.triggerPeriodicScheduledTasks();
                    context.runInMainThreadAndWait(() -> {});

                    assertThat(slotsTotalGauge.get().getValue())
                            .isEqualTo(2 * DEFAULT_NUM_SLOTS_PER_WORKER);
                    assertThat(slotsAvailableGauge.get().getValue())
                            .isEqualTo(2 * DEFAULT_NUM_SLOTS_PER_WORKER - 1);

                    final CompletableFuture<Boolean> unRegisterTaskManagerFuture =
                            new CompletableFuture<>();
                    context.runInMainThreadAndWait(
                            () ->
                                    unRegisterTaskManagerFuture.complete(
                                            context.getSlotManager()
                                                    .unregisterTaskManager(
                                                            executorConnection2.getInstanceID(),
                                                            TEST_EXCEPTION)));
                    assertThat(assertFutureCompleteAndReturn(unRegisterTaskManagerFuture)).isTrue();

                    // triggers the metric update task on the main thread and waits for the main
                    // thread to process queued up callbacks
                    scheduledExecutor.triggerPeriodicScheduledTasks();
                    context.runInMainThreadAndWait(() -> {});