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

Example 1 with SeqMapTask

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

the class DataflowWorkUnitClientTest method testCloudServiceCallSeqMapTaskStagePropagation.

@Test
public void testCloudServiceCallSeqMapTaskStagePropagation() throws Exception {
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    // Publish and acquire a seq map task work item, and verify we're now processing that stage.
    final String stageName = "test_stage_name";
    SeqMapTask seqMapTask = new SeqMapTask();
    seqMapTask.setStageName(stageName);
    WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
    workItem.setSeqMapTask(seqMapTask);
    when(request.execute()).thenReturn(generateMockResponse(workItem));
    assertEquals(Optional.of(workItem), client.getWorkItem());
    assertEquals(stageName, DataflowWorkerLoggingMDC.getStageName());
}