Examples with DataflowClient org.apache.beam.runners.dataflow.DataflowClient used on opensource projects

Example 1 with DataflowClient

use of org.apache.beam.runners.dataflow.DataflowClient in project beam by apache.

the class MonitoringUtilTest method testGetJobMessages.

@Test
public void testGetJobMessages() throws IOException {
    DataflowClient dataflowClient = mock(DataflowClient.class);
    ListJobMessagesResponse firstResponse = new ListJobMessagesResponse();
    firstResponse.setJobMessages(new ArrayList<JobMessage>());
    for (int i = 0; i < 100; ++i) {
        JobMessage message = new JobMessage();
        message.setId("message_" + i);
        message.setTime(TimeUtil.toCloudTime(new Instant(i)));
        firstResponse.getJobMessages().add(message);
    }
    String pageToken = "page_token";
    firstResponse.setNextPageToken(pageToken);
    ListJobMessagesResponse secondResponse = new ListJobMessagesResponse();
    secondResponse.setJobMessages(new ArrayList<JobMessage>());
    for (int i = 100; i < 150; ++i) {
        JobMessage message = new JobMessage();
        message.setId("message_" + i);
        message.setTime(TimeUtil.toCloudTime(new Instant(i)));
        secondResponse.getJobMessages().add(message);
    }
    when(dataflowClient.listJobMessages(JOB_ID, null)).thenReturn(firstResponse);
    when(dataflowClient.listJobMessages(JOB_ID, pageToken)).thenReturn(secondResponse);
    MonitoringUtil util = new MonitoringUtil(dataflowClient);
    List<JobMessage> messages = util.getJobMessages(JOB_ID, -1);
    assertEquals(150, messages.size());
}