Source:
private void runTest(TestCase testCase) throws Exception { AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime()); TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, testCase); HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource); RemoteTask remoteTask = httpRemoteTaskFactory.createRemoteTask( TEST_SESSION, new TaskId("test", 1, 2), new PrestoNode("node-id", URI.create("http://fake.invalid/"), new NodeVersion("version"), false), TaskTestUtils.PLAN_FRAGMENT, ImmutableMultimap.of(), createInitialEmptyOutputBuffers(OutputBuffers.BufferType.BROADCAST), new NodeTaskMap.PartitionedSplitCountTracker(i -> { }), true); testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo()); remoteTask.start(); CompletableFuture<Void> testComplete = new CompletableFuture<>(); asyncRun( IDLE_TIMEOUT.roundTo(MILLISECONDS), FAIL_TIMEOUT.roundTo(MILLISECONDS), lastActivityNanos, () -> testComplete.complete(null), (message, cause) -> testComplete.completeExceptionally(new AssertionError(message, cause))); testComplete.get(); httpRemoteTaskFactory.stop(); assertTrue(remoteTask.getTaskStatus().getState().isDone(), format("TaskStatus is not in a done state: %s", remoteTask.getTaskStatus())); assertTrue(remoteTask.getTaskInfo().getTaskStatus().getState().isDone(), format("TaskInfo is not in a done state: %s", remoteTask.getTaskInfo())); switch (testCase) { case TASK_MISMATCH: assertEquals(getOnlyElement(remoteTask.getTaskStatus().getFailures()).getErrorCode(), REMOTE_TASK_MISMATCH.toErrorCode()); break; case REJECTED_EXECUTION: assertEquals(getOnlyElement(remoteTask.getTaskStatus().getFailures()).getErrorCode(), REMOTE_TASK_ERROR.toErrorCode()); break; default: throw new UnsupportedOperationException(); } }

Target:
private void runTest(TestCase testCase) throws Exception { AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime()); TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, testCase); HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource); RemoteTask remoteTask = httpRemoteTaskFactory.createRemoteTask( TEST_SESSION, new TaskId("test", 1, 2), new PrestoNode("node-id", URI.create("http://fake.invalid/"), new NodeVersion("version"), false), TaskTestUtils.PLAN_FRAGMENT, ImmutableMultimap.of(), createInitialEmptyOutputBuffers(OutputBuffers.BufferType.BROADCAST), new NodeTaskMap.PartitionedSplitCountTracker(i -> { }), true); testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo()); remoteTask.start(); CompletableFuture<Void> testComplete = new CompletableFuture<>(); asyncRun( IDLE_TIMEOUT.roundTo(MILLISECONDS), FAIL_TIMEOUT.roundTo(MILLISECONDS), lastActivityNanos, () -> testComplete.complete(null), (message, cause) -> testComplete.completeExceptionally(new AssertionError(message, cause))); testComplete.get(); httpRemoteTaskFactory.stop(); assertTrue(remoteTask.getTaskStatus().getState().isDone(), format("TaskStatus is not in a done state: %s", remoteTask.getTaskStatus())); assertTrue(remoteTask.getTaskInfo().getTaskStatus().getState().isDone(), format("TaskInfo is not in a done state: %s", remoteTask.getTaskInfo())); ErrorCode actualErrorCode = getOnlyElement(remoteTask.getTaskStatus().getFailures()).getErrorCode(); switch (testCase) { case TASK_MISMATCH: case TASK_MISMATCH_WHEN_VERSION_IS_HIGH: assertEquals(actualErrorCode, REMOTE_TASK_MISMATCH.toErrorCode()); break; case REJECTED_EXECUTION: assertEquals(actualErrorCode, REMOTE_TASK_ERROR.toErrorCode()); break; default: throw new UnsupportedOperationException(); } }