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

Example 11 with WorkItemStatus

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

the class WorkItemStatusClientTest method populateCounterUpdatesEmptyOutputCounters.

@Test
public void populateCounterUpdatesEmptyOutputCounters() throws Exception {
    // When worker.getOutputCounters == null, there should be no counters.
    WorkItemStatus status = new WorkItemStatus();
    statusClient.setWorker(worker, executionContext);
    when(worker.extractMetricUpdates()).thenReturn(Collections.emptyList());
    statusClient.populateCounterUpdates(status);
    assertThat(status.getCounterUpdates(), hasSize(0));
}

Example 12 with WorkItemStatus

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

the class WorkItemStatusClientTest method populateCounterUpdatesWithOutputCounters.

@Test
public /**
 * Validates that an "internal" Counter is reported.
 */
void populateCounterUpdatesWithOutputCounters() throws Exception {
    final CounterUpdate counter = new CounterUpdate().setNameAndKind(new NameAndKind().setName("some-counter").setKind(Kind.SUM.toString())).setCumulative(true).setInteger(DataflowCounterUpdateExtractor.longToSplitInt(42));
    CounterSet counterSet = new CounterSet();
    counterSet.intSum(CounterName.named("some-counter")).addValue(42);
    WorkItemStatus status = new WorkItemStatus();
    when(worker.getOutputCounters()).thenReturn(counterSet);
    when(worker.extractMetricUpdates()).thenReturn(Collections.emptyList());
    when(worker.extractMetricUpdates()).thenReturn(Collections.emptyList());
    statusClient.setWorker(worker, executionContext);
    statusClient.populateCounterUpdates(status);
    assertThat(status.getCounterUpdates(), containsInAnyOrder(counter));
}

Example 13 with WorkItemStatus

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

the class WorkItemStatusClientTest method populateSplitResultCustomReader.

@Test
public void populateSplitResultCustomReader() throws Exception {
    WorkItemStatus status = new WorkItemStatus();
    statusClient.setWorker(worker, executionContext);
    BoundedSource<Integer> primary = new DummyBoundedSource(5);
    BoundedSource<Integer> residual = new DummyBoundedSource(10);
    BoundedSourceSplit<Integer> split = new BoundedSourceSplit<>(primary, residual);
    statusClient.populateSplitResult(status, split);
    assertThat(status.getDynamicSourceSplit(), equalTo(WorkerCustomSources.toSourceSplit(split)));
    assertThat(status.getStopPosition(), nullValue());
}

Example 14 with WorkItemStatus

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

the class WorkItemStatusClientTest method populateCounterUpdatesWithMetricsAndCounters.

/**
 * Validates that Beam Metrics and "internal" Counters are merged in the update.
 */
@Test
public void populateCounterUpdatesWithMetricsAndCounters() throws Exception {
    final CounterUpdate expectedCounter = new CounterUpdate().setNameAndKind(new NameAndKind().setName("some-counter").setKind(Kind.SUM.toString())).setCumulative(true).setInteger(DataflowCounterUpdateExtractor.longToSplitInt(42));
    CounterSet counterSet = new CounterSet();
    counterSet.intSum(CounterName.named("some-counter")).addValue(42);
    final CounterUpdate expectedMetric = new CounterUpdate().setStructuredNameAndMetadata(new CounterStructuredNameAndMetadata().setName(new CounterStructuredName().setOrigin("USER").setOriginNamespace("namespace").setName("some-counter").setOriginalStepName("step")).setMetadata(new CounterMetadata().setKind(Kind.SUM.toString()))).setCumulative(true).setInteger(DataflowCounterUpdateExtractor.longToSplitInt(42));
    MetricsContainerImpl metricsContainer = new MetricsContainerImpl("step");
    BatchModeExecutionContext context = mock(BatchModeExecutionContext.class);
    when(context.extractMetricUpdates(anyBoolean())).thenReturn(ImmutableList.of(expectedMetric));
    when(context.extractMsecCounters(anyBoolean())).thenReturn(Collections.emptyList());
    CounterCell counter = metricsContainer.getCounter(MetricName.named("namespace", "some-counter"));
    counter.inc(1);
    counter.inc(41);
    counter.inc(1);
    counter.inc(-1);
    WorkItemStatus status = new WorkItemStatus();
    when(worker.getOutputCounters()).thenReturn(counterSet);
    when(worker.extractMetricUpdates()).thenReturn(Collections.emptyList());
    statusClient.setWorker(worker, context);
    statusClient.populateCounterUpdates(status);
    assertThat(status.getCounterUpdates(), containsInAnyOrder(expectedCounter, expectedMetric));
}

Example 15 with WorkItemStatus

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

the class WorkItemStatusClientTest method reportSuccessWithSourceOperation.

@Test
public void reportSuccessWithSourceOperation() throws IOException {
    SourceOperationExecutor sourceWorker = mock(SourceOperationExecutor.class);
    when(sourceWorker.extractMetricUpdates()).thenReturn(Collections.emptyList());
    statusClient.setWorker(sourceWorker, 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());
}