Examples with TestDataflowExecutionState org.apache.beam.runners.dataflow.worker.TestOperationContext.TestDataflowExecutionState used on opensource projects

Example 6 with TestDataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.TestOperationContext.TestDataflowExecutionState in project beam by apache.

the class MapTaskExecutorTest method testGetMetricContainers.

@Test
@SuppressWarnings("unchecked")
public /**
 * This test makes sure that any metrics reported within an operation are part of the metric
 * containers returned by {@link getMetricContainers}.
 */
void testGetMetricContainers() throws Exception {
    ExecutionStateTracker stateTracker = new DataflowExecutionStateTracker(ExecutionStateSampler.newForTest(), new TestDataflowExecutionState(NameContext.forStage("testStage"), "other", null, /* requestingStepName */
    null, /* sideInputIndex */
    null, /* metricsContainer */
    NoopProfileScope.NOOP), new CounterSet(), PipelineOptionsFactory.create(), "test-work-item-id");
    final String o1 = "o1";
    TestOperationContext context1 = createContext(o1, stateTracker);
    final String o2 = "o2";
    TestOperationContext context2 = createContext(o2, stateTracker);
    final String o3 = "o3";
    TestOperationContext context3 = createContext(o3, stateTracker);
    List<Operation> operations = Arrays.asList(new Operation(new OutputReceiver[] {}, context1) {

        @Override
        public void start() throws Exception {
            super.start();
            try (Closeable scope = context.enterStart()) {
                Metrics.counter("TestMetric", "MetricCounter").inc(1L);
            }
        }
    }, new Operation(new OutputReceiver[] {}, context2) {

        @Override
        public void start() throws Exception {
            super.start();
            try (Closeable scope = context.enterStart()) {
                Metrics.counter("TestMetric", "MetricCounter").inc(2L);
            }
        }
    }, new Operation(new OutputReceiver[] {}, context3) {

        @Override
        public void start() throws Exception {
            super.start();
            try (Closeable scope = context.enterStart()) {
                Metrics.counter("TestMetric", "MetricCounter").inc(3L);
            }
        }
    });
    try (MapTaskExecutor executor = new MapTaskExecutor(operations, counterSet, stateTracker)) {
        // Call execute so that we run all the counters
        executor.execute();
        assertThat(context1.metricsContainer().getUpdates().counterUpdates(), contains(metricUpdate("TestMetric", "MetricCounter", o1, 1L)));
        assertThat(context2.metricsContainer().getUpdates().counterUpdates(), contains(metricUpdate("TestMetric", "MetricCounter", o2, 2L)));
        assertThat(context3.metricsContainer().getUpdates().counterUpdates(), contains(metricUpdate("TestMetric", "MetricCounter", o3, 3L)));
    }
}

Example 7 with TestDataflowExecutionState

use of org.apache.beam.runners.dataflow.worker.TestOperationContext.TestDataflowExecutionState in project beam by apache.

the class DataflowWorkerLoggingHandlerTest method testWithAllValuesInMDC.

@Test
public void testWithAllValuesInMDC() throws IOException {
    DataflowExecutionState state = new TestDataflowExecutionState(nameContextForTest(), "activity");
    tracker.enterState(state);
    String testJobId = "testJobId";
    String testStage = "testStage";
    String testWorkerId = "testWorkerId";
    String testWorkId = "testWorkId";
    DataflowWorkerLoggingMDC.setJobId(testJobId);
    DataflowWorkerLoggingMDC.setStageName(testStage);
    DataflowWorkerLoggingMDC.setWorkerId(testWorkerId);
    DataflowWorkerLoggingMDC.setWorkId(testWorkId);
    createLogRecord("test.message", null);
    assertEquals(String.format("{\"timestamp\":{\"seconds\":0,\"nanos\":1000000},\"severity\":\"INFO\"," + "\"message\":\"test.message\",\"thread\":\"2\",\"job\":\"%s\"," + "\"stage\":\"%s\",\"step\":\"%s\",\"worker\":\"%s\"," + "\"work\":\"%s\",\"logger\":\"LoggerName\"}" + System.lineSeparator(), testJobId, testStage, NameContextsForTests.USER_NAME, testWorkerId, testWorkId), createJson(createLogRecord("test.message", null)));
}