Examples with WorkItemStatus com.google.api.services.dataflow.model.WorkItemStatus used on opensource projects

Example 16 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class WorkItemStatusClientTest method populateMetricUpdatesStateSamplerInfo.

@Test
public void populateMetricUpdatesStateSamplerInfo() throws Exception {
    // When executionContext.getExecutionStateTracker() returns non-null, we get one metric update.
    WorkItemStatus status = new WorkItemStatus();
    BatchModeExecutionContext executionContext = mock(BatchModeExecutionContext.class);
    ExecutionStateTracker executionStateTracker = mock(ExecutionStateTracker.class);
    ExecutionState executionState = mock(ExecutionState.class);
    when(executionState.getDescription()).thenReturn("stageName-systemName-some-state");
    when(executionContext.getExecutionStateTracker()).thenReturn(executionStateTracker);
    when(executionStateTracker.getMillisSinceLastTransition()).thenReturn(20L);
    when(executionStateTracker.getNumTransitions()).thenReturn(10L);
    when(executionStateTracker.getCurrentState()).thenReturn(executionState);
    statusClient.setWorker(worker, executionContext);
    statusClient.populateMetricUpdates(status);
    assertThat(status.getMetricUpdates(), hasSize(1));
    MetricUpdate update = status.getMetricUpdates().get(0);
    assertThat(update.getName().getName(), equalTo("state-sampler"));
    assertThat(update.getKind(), equalTo("internal"));
    Map<String, Object> samplerMetrics = (Map<String, Object>) update.getInternal();
    assertThat(samplerMetrics, hasEntry("last-state-name", "stageName-systemName-some-state"));
    assertThat(samplerMetrics, hasEntry("num-transitions", 10L));
    assertThat(samplerMetrics, hasEntry("last-state-duration-ms", 20L));
}

Example 17 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class WorkItemStatusClientTest method populateProgressNull.

@Test
public void populateProgressNull() throws Exception {
    WorkItemStatus status = new WorkItemStatus();
    statusClient.setWorker(worker, executionContext);
    statusClient.populateProgress(status);
    assertThat(status.getReportedProgress(), nullValue());
}

Example 18 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class WorkItemStatusClientTest method reportUpdate.

@Test
public void reportUpdate() throws Exception {
    when(worker.extractMetricUpdates()).thenReturn(Collections.emptyList());
    statusClient.setWorker(worker, executionContext);
    statusClient.reportUpdate(null, LEASE_DURATION);
    verify(workUnitClient).reportWorkItemStatus(statusCaptor.capture());
    WorkItemStatus workStatus = statusCaptor.getValue();
    assertThat(workStatus.getCompleted(), equalTo(false));
}

Example 19 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class WorkItemStatusClientTest method reportError.

/**
 * Reporting an error before setWorker has been called should work.
 */
@Test
public void reportError() throws IOException {
    RuntimeException error = new RuntimeException();
    error.fillInStackTrace();
    statusClient.reportError(error);
    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(), hasSize(1));
    Status status = workStatus.getErrors().get(0);
    assertThat(status.getCode(), equalTo(2));
    assertThat(status.getMessage(), containsString("WorkItemStatusClientTest"));
}

Example 20 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class WorkItemStatusClient method reportSuccess.

/**
 * Return the {@link WorkItemServiceState} resulting from sending a success completion status.
 */
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) {
        // TODO: Find out a generic way for the DataflowWorkExecutor to report work-specific results
        // into the work update.
        SourceOperationResponse response = ((SourceOperationExecutor) worker).getResponse();
        if (response != null) {
            status.setSourceOperationResponse(response);
        }
    }
    LOG.info("Success processing work item {}", uniqueWorkId());
    return execute(status);
}