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

Example 16 with WorkItem

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

the class DataflowWorkUnitClientTest method testCloudServiceCallNoWorkPresent.

@Test
public void testCloudServiceCallNoWorkPresent() throws Exception {
    // If there's no work the service should return an empty work item.
    WorkItem workItem = new WorkItem();
    when(request.execute()).thenReturn(generateMockResponse(workItem));
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    assertEquals(Optional.absent(), client.getWorkItem());
    LeaseWorkItemRequest actualRequest = Transport.getJsonFactory().fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
    assertEquals(WORKER_ID, actualRequest.getWorkerId());
    assertEquals(ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities());
    assertEquals(ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes());
}

Example 17 with WorkItem

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

the class DataflowWorkUnitClientTest method testCloudServiceCall.

@Test
public void testCloudServiceCall() throws Exception {
    WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
    when(request.execute()).thenReturn(generateMockResponse(workItem));
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    assertEquals(Optional.of(workItem), client.getWorkItem());
    LeaseWorkItemRequest actualRequest = Transport.getJsonFactory().fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
    assertEquals(WORKER_ID, actualRequest.getWorkerId());
    assertEquals(ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities());
    assertEquals(ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes());
    assertEquals("1234", DataflowWorkerLoggingMDC.getWorkId());
}

Example 18 with WorkItem

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

the class DataflowWorkUnitClientTest method createWorkItem.

private WorkItem createWorkItem(String projectId, String jobId) {
    WorkItem workItem = new WorkItem();
    workItem.setFactory(Transport.getJsonFactory());
    workItem.setProjectId(projectId);
    workItem.setJobId(jobId);
    // We need to set a work id because otherwise the client will treat the response as
    // indicating no work is available.
    workItem.setId(1234L);
    return workItem;
}

Example 19 with WorkItem

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

the class DataflowWorkProgressUpdaterTest method initMocksAndWorkflowServiceAndWorkerAndWork.

@Before
public void initMocksAndWorkflowServiceAndWorkerAndWork() {
    MockitoAnnotations.initMocks(this);
    startTime = 0L;
    clock = new FixedClock(startTime);
    executor = new StubbedExecutor(clock);
    options = PipelineOptionsFactory.create().as(DataflowPipelineOptions.class);
    options.setHotKeyLoggingEnabled(true);
    WorkItem workItem = new WorkItem();
    workItem.setProjectId(PROJECT_ID);
    workItem.setJobId(JOB_ID);
    workItem.setId(WORK_ID);
    workItem.setLeaseExpireTime(toCloudTime(new Instant(clock.currentTimeMillis() + 1000)));
    workItem.setReportStatusInterval(toCloudDuration(Duration.millis(300)));
    workItem.setInitialReportIndex(1L);
    when(workItemStatusClient.getExecutionContext()).thenReturn(context);
    when(context.getKey()).thenReturn(HOT_KEY);
    progressUpdater = new DataflowWorkProgressUpdater(workItemStatusClient, workItem, worker, executor.getExecutor(), clock, hotKeyLogger, options) {

        // Shorten reporting interval boundaries for faster testing.
        @Override
        protected long getMinReportingInterval() {
            return 100;
        }

        @Override
        protected long getLeaseRenewalLatencyMargin() {
            return 150;
        }
    };
}