Source:
public void testExecutionDeploymentReconciliation() throws Exception { JobMasterBuilder.TestingOnCompletionActions onCompletionActions = new JobMasterBuilder.TestingOnCompletionActions(); JobMaster jobMaster = createAndStartJobMaster(onCompletionActions); JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class); rpcService.registerGateway(jobMasterGateway.getAddress(), jobMasterGateway); final JobMasterPartitionReleaseTest.AllocationIdsResourceManagerGateway resourceManagerGateway = createResourceManagerGateway(); final CompletableFuture<Void> taskSubmissionFuture = new CompletableFuture<>(); final CompletableFuture<ExecutionAttemptID> taskCancellationFuture = new CompletableFuture<>(); TaskExecutorGateway taskExecutorGateway = createTaskExecutorGateway(taskSubmissionFuture, taskCancellationFuture); LocalUnresolvedTaskManagerLocation localUnresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation(); registerTaskExecutorAndOfferSlots(resourceManagerGateway, jobMasterGateway, taskExecutorGateway, localUnresolvedTaskManagerLocation); taskSubmissionFuture.get(); assertFalse(taskCancellationFuture.isDone()); ExecutionAttemptID unknownDeployment = new ExecutionAttemptID(); jobMasterGateway.heartbeatFromTaskManager(localUnresolvedTaskManagerLocation.getResourceID(), new TaskExecutorToJobManagerHeartbeatPayload( new AccumulatorReport(Collections.emptyList()), new ExecutionDeploymentReport(Collections.singleton(unknownDeployment)) )); assertThat(taskCancellationFuture.get(), is(unknownDeployment)); assertThat(onCompletionActions.getJobReachedGloballyTerminalStateFuture().get().getState(), is(JobStatus.FAILED)); }

Target:
public void testExecutionDeploymentReconciliation() throws Exception { JobMasterBuilder.TestingOnCompletionActions onCompletionActions = new JobMasterBuilder.TestingOnCompletionActions(); final CompletableFuture<Void> taskDeploymentFuture = new CompletableFuture<>(); JobMaster jobMaster = createAndStartJobMaster(onCompletionActions, taskDeploymentFuture); JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class); RPC_SERVICE_RESOURCE.getTestingRpcService().registerGateway(jobMasterGateway.getAddress(), jobMasterGateway); final AllocationIdsExposingResourceManagerGateway resourceManagerGateway = createResourceManagerGateway(); final CompletableFuture<ExecutionAttemptID> taskCancellationFuture = new CompletableFuture<>(); TaskExecutorGateway taskExecutorGateway = createTaskExecutorGateway(taskCancellationFuture); LocalUnresolvedTaskManagerLocation localUnresolvedTaskManagerLocation = new LocalUnresolvedTaskManagerLocation(); registerTaskExecutorAndOfferSlots(resourceManagerGateway, jobMasterGateway, taskExecutorGateway, localUnresolvedTaskManagerLocation); taskDeploymentFuture.get(); assertFalse(taskCancellationFuture.isDone()); ExecutionAttemptID unknownDeployment = new ExecutionAttemptID(); jobMasterGateway.heartbeatFromTaskManager(localUnresolvedTaskManagerLocation.getResourceID(), new TaskExecutorToJobManagerHeartbeatPayload( new AccumulatorReport(Collections.emptyList()), new ExecutionDeploymentReport(Collections.singleton(unknownDeployment)) )); assertThat(taskCancellationFuture.get(), is(unknownDeployment)); assertThat(onCompletionActions.getJobReachedGloballyTerminalStateFuture().get().getState(), is(JobStatus.FAILED)); }