Source:
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception { final int parallelism = 5; final Time batchSlotTimeout = Time.milliseconds(5L); final JobGraph jobGraph = createJobGraph(parallelism); jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES_WITH_BATCH_SLOT_REQUEST); try (final SlotPoolImpl slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) { final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism); TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder() .setSubmitTaskConsumer( (tdd, ignored) -> { submittedTasksQueue.offer(tdd.getExecutionAttemptId()); return CompletableFuture.completedFuture(Acknowledge.get()); }) .createTestingTaskExecutorGateway(); SlotPoolUtils.offerSlots( slotPool, mainThreadExecutor, Collections.singletonList(ResourceProfile.ANY), new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate())); final PhysicalSlotProvider slotProvider = createSlotProvider(slotPool, mainThreadExecutor); final SchedulerNG scheduler = createScheduler(jobGraph, slotProvider, batchSlotTimeout); final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener(); scheduler.registerJobStatusListener(jobStatusListener); startScheduling(scheduler, mainThreadExecutor); Thread.sleep(batchSlotTimeout.toMilliseconds()); final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture(); for (int i = 0; i < parallelism; i++) { final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync(CheckedSupplier.unchecked(submittedTasksQueue::take)); CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join(); if (submittedTaskFuture.isDone()) { finishExecution(submittedTaskFuture.get(), scheduler, mainThreadExecutor); } else { fail(String.format("Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get())); } } assertThat(terminationFuture.get(), is(JobStatus.FINISHED)); } }

Target:
public void testSchedulingOfJobWithFewerSlotsThanParallelism() throws Exception { final int parallelism = 5; final Time batchSlotTimeout = Time.milliseconds(5L); final JobGraph jobGraph = createJobGraph(parallelism); jobGraph.setScheduleMode(ScheduleMode.LAZY_FROM_SOURCES_WITH_BATCH_SLOT_REQUEST); try (final SlotPoolImpl slotPool = createSlotPool(mainThreadExecutor, batchSlotTimeout)) { final ArrayBlockingQueue<ExecutionAttemptID> submittedTasksQueue = new ArrayBlockingQueue<>(parallelism); TestingTaskExecutorGateway testingTaskExecutorGateway = new TestingTaskExecutorGatewayBuilder() .setSubmitTaskConsumer( (tdd, ignored) -> { submittedTasksQueue.offer(tdd.getExecutionAttemptId()); return CompletableFuture.completedFuture(Acknowledge.get()); }) .createTestingTaskExecutorGateway(); SlotPoolUtils.offerSlots( slotPool, mainThreadExecutor, Collections.singletonList(ResourceProfile.ANY), new RpcTaskManagerGateway(testingTaskExecutorGateway, JobMasterId.generate())); final PhysicalSlotProvider slotProvider = new PhysicalSlotProviderImpl( LocationPreferenceSlotSelectionStrategy.createDefault(), slotPool); final SchedulerNG scheduler = createScheduler(jobGraph, slotProvider, batchSlotTimeout); final GloballyTerminalJobStatusListener jobStatusListener = new GloballyTerminalJobStatusListener(); scheduler.registerJobStatusListener(jobStatusListener); startScheduling(scheduler, mainThreadExecutor); Thread.sleep(batchSlotTimeout.toMilliseconds()); final CompletableFuture<JobStatus> terminationFuture = jobStatusListener.getTerminationFuture(); for (int i = 0; i < parallelism; i++) { final CompletableFuture<ExecutionAttemptID> submittedTaskFuture = CompletableFuture.supplyAsync( CheckedSupplier.unchecked(submittedTasksQueue::take)); CompletableFuture.anyOf(submittedTaskFuture, terminationFuture).join(); if (submittedTaskFuture.isDone()) { finishExecution(submittedTaskFuture.get(), scheduler, mainThreadExecutor); } else { fail( String.format( "Job reached a globally terminal state %s before all executions were finished.", terminationFuture.get())); } } assertThat(terminationFuture.get(), is(JobStatus.FINISHED)); } }