@VisibleForTesting
  boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) throws IOException {
    LOG.debug("Executing: {}", workItem);

    DataflowWorkExecutor worker = null;
    try {
      // Populate PipelineOptions with data from work unit.
      options.setProject(workItem.getProjectId());

      final String stageName;
      if (workItem.getMapTask() != null) {
        stageName = workItem.getMapTask().getStageName();
      } else if (workItem.getSourceOperationTask() != null) {
        stageName = workItem.getSourceOperationTask().getStageName();
      } else {
        throw new RuntimeException("Unknown kind of work item: " + workItem);
      }

      CounterSet counterSet = new CounterSet();
      BatchModeExecutionContext executionContext =
          BatchModeExecutionContext.create(
              counterSet,
              sideInputDataCache,
              sideInputWeakReferenceCache,
              readerRegistry,
              options,
              stageName,
              String.valueOf(workItem.getId()));

      if (workItem.getMapTask() != null) {
        MutableNetwork<Node, Edge> network = mapTaskToNetwork.apply(workItem.getMapTask());
        if (LOG.isDebugEnabled()) {
          LOG.debug("Network as Graphviz .dot: {}", Networks.toDot(network));
        }

        worker =
            mapTaskExecutorFactory.create(
                network,
                options,
                stageName,
                readerRegistry,
                sinkRegistry,
                executionContext,
                counterSet,
                idGenerator);
      } else if (workItem.getSourceOperationTask() != null) {
        worker =
            SourceOperationExecutorFactory.create(
                options,
                workItem.getSourceOperationTask(),
                counterSet,
                executionContext,
                stageName);
      } else {
        throw new IllegalStateException("Work Item was neither a MapTask nor a SourceOperation");
      }
      workItemStatusClient.setWorker(worker, executionContext);

      DataflowWorkProgressUpdater progressUpdater =
          new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, options);
      executeWork(worker, progressUpdater);
      workItemStatusClient.reportSuccess();
      return true;

    } catch (Throwable e) {
      workItemStatusClient.reportError(e);
      return false;

    } finally {
      if (worker != null) {
        try {
          worker.close();
        } catch (Exception exn) {
          LOG.warn(
              "Uncaught exception while closing worker. All work has already committed or "
                  + "been marked for retry.",
              exn);
        }
      }
    }
  }
@Test
  public void testWhenProcessingWorkUnitFailsWeReportStatus() throws Exception {
    BatchDataflowWorker worker =
        new BatchDataflowWorker(
            mockWorkUnitClient, IntrinsicMapTaskExecutorFactory.defaultFactory(), options);
    // In practice this value is always 1, but for the sake of testing send a different value.
    long initialReportIndex = 4L;
    WorkItem workItem =
        new WorkItem()
            .setId(1L)
            .setJobId("Expected to fail the job")
            .setInitialReportIndex(initialReportIndex);
    WorkItemStatusClient workItemStatusClient = mock(WorkItemStatusClient.class);

    worker.doWork(workItem, workItemStatusClient);
    ArgumentCaptor<Throwable> errorCaptor = ArgumentCaptor.forClass(Throwable.class);
    verify(workItemStatusClient).reportError(errorCaptor.capture());

    Throwable error = errorCaptor.getValue();
    assertThat(error, notNullValue());
    assertThat(error.getMessage(), equalTo("Unknown kind of work item: " + workItem));
  }