WorkItemStatusClient { public synchronized WorkItemServiceState reportSuccess() throws IOException { checkState(!finalStateSent, "cannot reportSuccess after sending a final state"); checkState(worker != null, "setWorker should be called before reportSuccess"); WorkItemStatus status = createStatusUpdate(true); if (worker instanceof SourceOperationExecutor) { SourceOperationResponse response = ((SourceOperationExecutor) worker).getResponse(); if (response != null) { status.setSourceOperationResponse(response); } } LOG.info("Success processing work item {}", uniqueWorkId()); return execute(status); } WorkItemStatusClient(WorkUnitClient workUnitClient, WorkItem workItem); String uniqueWorkId(); synchronized void setWorker(
      DataflowWorkExecutor worker, BatchModeExecutionContext executionContext); synchronized WorkItemServiceState reportError(Throwable e); synchronized WorkItemServiceState reportSuccess(); synchronized WorkItemServiceState reportUpdate(
      @Nullable DynamicSplitResult dynamicSplitResult, Duration requestedLeaseDuration); Iterable<CounterUpdate> extractMsecCounters(boolean isFinalUpdate); long extractThrottleTime();  }
@Test public void reportSuccessBeforeSetWorker() throws IOException { thrown.expect(IllegalStateException.class); thrown.expectMessage("setWorker"); thrown.expectMessage("reportSuccess"); statusClient.reportSuccess(); }
@Test public void reportSuccess() throws IOException { when(worker.extractMetricUpdates()).thenReturn(Collections.emptyList()); statusClient.setWorker(worker, executionContext); statusClient.reportSuccess(); verify(workUnitClient).reportWorkItemStatus(statusCaptor.capture()); WorkItemStatus workStatus = statusCaptor.getValue(); assertThat(workStatus.getWorkItemId(), equalTo(Long.toString(WORK_ID))); assertThat(workStatus.getCompleted(), equalTo(true)); assertThat(workStatus.getReportIndex(), equalTo(INITIAL_REPORT_INDEX)); assertThat(workStatus.getErrors(), nullValue()); }