public static String getGcloudCancelCommand(DataflowPipelineOptions options, String jobId) {

    // If using a different Dataflow API than default, prefix command with an API override.
    String dataflowApiOverridePrefix = "";
    String apiUrl = options.getDataflowClient().getBaseUrl();
    if (!apiUrl.equals(Dataflow.DEFAULT_BASE_URL)) {
      dataflowApiOverridePrefix = String.format("%s=%s ", ENDPOINT_OVERRIDE_ENV_VAR, apiUrl);
    }

    // Assemble cancel command from optional prefix and project/job parameters.
    return String.format(
        "%s%s jobs --project=%s cancel --region=%s %s",
        dataflowApiOverridePrefix,
        GCLOUD_DATAFLOW_PREFIX,
        options.getProject(),
        options.getRegion(),
        jobId);
  }
 

/** Returns a Google Cloud Dataflow client builder. */
public static Dataflow.Builder newDataflowClient(DataflowPipelineOptions options) {
  String servicePath = options.getDataflowEndpoint();
  ApiComponents components;
  if (servicePath.contains("://")) {
    components = apiComponentsFromUrl(servicePath);
  } else {
    components = new ApiComponents(options.getApiRootUrl(), servicePath);
  }

  return new Dataflow.Builder(
          getTransport(),
          getJsonFactory(),
          chainHttpRequestInitializer(
              options.getGcpCredential(),
              // Do not log 404. It clutters the output and is possibly even required by the
              // caller.
              new RetryHttpRequestInitializer(ImmutableList.of(404))))
      .setApplicationName(options.getAppName())
      .setRootUrl(components.rootUrl)
      .setServicePath(components.servicePath)
      .setGoogleClientRequestInitializer(options.getGoogleApiTrace());
}
 

@Test
public void testWaitToFinishMessagesFail() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_" + State.DONE.name());
  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  MonitoringUtil.JobMessagesHandler jobHandler = mock(MonitoringUtil.JobMessagesHandler.class);
  Dataflow.Projects.Locations.Jobs.Messages mockMessages =
      mock(Dataflow.Projects.Locations.Jobs.Messages.class);
  Messages.List listRequest = mock(Dataflow.Projects.Locations.Jobs.Messages.List.class);
  when(mockJobs.messages()).thenReturn(mockMessages);
  when(mockMessages.list(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(listRequest);
  when(listRequest.setPageToken(eq((String) null))).thenReturn(listRequest);
  when(listRequest.execute()).thenThrow(SocketTimeoutException.class);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());

  State state =
      job.waitUntilFinish(Duration.standardMinutes(5), jobHandler, fastClock, fastClock);
  assertEquals(null, state);
}
 

public State mockWaitToFinishInState(State state) throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_" + state.name());
  if (state == State.UPDATED) {
    statusResponse.setReplacedByJobId(REPLACEMENT_JOB_ID);
  }

  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());

  return job.waitUntilFinish(Duration.standardMinutes(1), null, fastClock, fastClock);
}
 

@Test
public void testWaitToFinishFail() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenThrow(IOException.class);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());

  long startTime = fastClock.nanoTime();
  State state = job.waitUntilFinish(Duration.standardMinutes(5), null, fastClock, fastClock);
  assertEquals(null, state);
  long timeDiff = TimeUnit.NANOSECONDS.toMillis(fastClock.nanoTime() - startTime);
  checkValidInterval(
      DataflowPipelineJob.MESSAGES_POLLING_INTERVAL,
      DataflowPipelineJob.MESSAGES_POLLING_RETRIES,
      timeDiff);
}
 

@Test
public void testWaitToFinishTimeFail() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenThrow(IOException.class);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());
  long startTime = fastClock.nanoTime();
  State state = job.waitUntilFinish(Duration.millis(4), null, fastClock, fastClock);
  assertEquals(null, state);
  long timeDiff = TimeUnit.NANOSECONDS.toMillis(fastClock.nanoTime() - startTime);
  // Should only have slept for the 4 ms allowed.
  assertEquals(4L, timeDiff);
}
 

@Test
public void testCumulativeTimeOverflow() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_RUNNING");
  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  FastNanoClockAndFuzzySleeper clock = new FastNanoClockAndFuzzySleeper();

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());
  long startTime = clock.nanoTime();
  State state = job.waitUntilFinish(Duration.millis(4), null, clock, clock);
  assertEquals(null, state);
  long timeDiff = TimeUnit.NANOSECONDS.toMillis(clock.nanoTime() - startTime);
  // Should only have slept for the 4 ms allowed.
  assertThat(timeDiff, lessThanOrEqualTo(4L));
}
 

@Test
public void testGetStateReturnsServiceState() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_" + State.RUNNING.name());

  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());

  assertEquals(
      State.RUNNING,
      job.getStateWithRetriesOrUnknownOnException(
          BackOffAdapter.toGcpBackOff(DataflowPipelineJob.STATUS_BACKOFF_FACTORY.backoff()),
          fastClock));
}
 

@Test
public void testGetStateWithRetriesPassesExceptionThrough() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenThrow(IOException.class);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());

  long startTime = fastClock.nanoTime();
  thrown.expect(IOException.class);
  job.getStateWithRetries(
      BackOffAdapter.toGcpBackOff(DataflowPipelineJob.STATUS_BACKOFF_FACTORY.backoff()),
      fastClock);
}
 

@Test
public void testGetStateNoThrowWithExceptionReturnsUnknown() throws Exception {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest);
  when(statusRequest.execute()).thenThrow(IOException.class);

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of());

  long startTime = fastClock.nanoTime();
  assertEquals(
      State.UNKNOWN,
      job.getStateWithRetriesOrUnknownOnException(
          BackOffAdapter.toGcpBackOff(DataflowPipelineJob.STATUS_BACKOFF_FACTORY.backoff()),
          fastClock));
  long timeDiff = TimeUnit.NANOSECONDS.toMillis(fastClock.nanoTime() - startTime);
  checkValidInterval(
      DataflowPipelineJob.STATUS_POLLING_INTERVAL,
      DataflowPipelineJob.STATUS_POLLING_RETRIES,
      timeDiff);
}
 

@Test
public void testCancelUnterminatedJobThatSucceeds() throws IOException {
  Dataflow.Projects.Locations.Jobs.Update update =
      mock(Dataflow.Projects.Locations.Jobs.Update.class);
  when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class)))
      .thenReturn(update);
  when(update.execute()).thenReturn(new Job().setCurrentState("JOB_STATE_CANCELLED"));

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null);

  assertEquals(State.CANCELLED, job.cancel());
  Job content = new Job();
  content.setProjectId(PROJECT_ID);
  content.setId(JOB_ID);
  content.setRequestedState("JOB_STATE_CANCELLED");
  verify(mockJobs).update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), eq(content));
  verifyNoMoreInteractions(mockJobs);
}
 

@Test
public void testCancelUnterminatedJobThatFails() throws IOException {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_RUNNING");
  when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  Dataflow.Projects.Locations.Jobs.Update update =
      mock(Dataflow.Projects.Locations.Jobs.Update.class);
  when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class)))
      .thenReturn(update);
  when(update.execute()).thenThrow(new IOException("Some random IOException"));

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null);

  thrown.expect(IOException.class);
  thrown.expectMessage(
      "Failed to cancel job in state RUNNING, "
          + "please go to the Developers Console to cancel it manually:");
  job.cancel();
}
 

/**
 * Test that {@link DataflowPipelineJob#cancel} doesn't throw if the Dataflow service returns
 * non-terminal state even though the cancel API call failed, which can happen in practice.
 *
 * <p>TODO: delete this code if the API calls become consistent.
 */
@Test
public void testCancelTerminatedJobWithStaleState() throws IOException {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_RUNNING");
  when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  Dataflow.Projects.Locations.Jobs.Update update =
      mock(Dataflow.Projects.Locations.Jobs.Update.class);
  when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class)))
      .thenReturn(update);
  when(update.execute()).thenThrow(new IOException("Job has terminated in state SUCCESS"));

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null);
  State returned = job.cancel();
  assertThat(returned, equalTo(State.RUNNING));
  expectedLogs.verifyWarn("Cancel failed because job is already terminated.");
}
 

private Dataflow buildMockDataflow() throws Exception {
  Dataflow mockDataflowClient = mock(Dataflow.class);
  Dataflow.Projects mockProjects = mock(Dataflow.Projects.class);
  Dataflow.Projects.Locations mockRegion = mock(Dataflow.Projects.Locations.class);
  Dataflow.Projects.Locations.Jobs mockJobs = mock(Dataflow.Projects.Locations.Jobs.class);
  Dataflow.Projects.Locations.Jobs.Debug mockDebug =
      mock(Dataflow.Projects.Locations.Jobs.Debug.class);
  Dataflow.Projects.Locations.Jobs.Debug.GetConfig mockGetConfig =
      mock(Dataflow.Projects.Locations.Jobs.Debug.GetConfig.class);
  when(mockDataflowClient.projects()).thenReturn(mockProjects);
  when(mockProjects.locations()).thenReturn(mockRegion);
  when(mockRegion.jobs()).thenReturn(mockJobs);
  when(mockJobs.debug()).thenReturn(mockDebug);
  when(mockDebug.getConfig(
          eq(PROJECT_ID), eq(REGION), eq(JOB_ID), isA(GetDebugConfigRequest.class)))
      .thenReturn(mockGetConfig);
  when(mockDebug.sendCapture(
          eq(PROJECT_ID), eq(REGION), eq(JOB_ID), isA(SendDebugCaptureRequest.class)))
      .thenReturn(mockSendCapture);

  when(mockGetConfig.execute()).thenReturn(fakeGetConfigResponse);

  return mockDataflowClient;
}
 

@Inject
GenerateSpec11ReportAction(
    @Config("projectId") String projectId,
    @Config("apacheBeamBucketUrl") String beamBucketUrl,
    @Config("spec11TemplateUrl") String spec11TemplateUrl,
    @Config("defaultJobZone") String jobZone,
    @Key("safeBrowsingAPIKey") String apiKey,
    @Parameter(PARAM_DATE) LocalDate date,
    Response response,
    Dataflow dataflow) {
  this.projectId = projectId;
  this.beamBucketUrl = beamBucketUrl;
  this.spec11TemplateUrl = spec11TemplateUrl;
  this.jobZone = jobZone;
  this.apiKey = apiKey;
  this.date = date;
  this.response = response;
  this.dataflow = dataflow;
}
 

@Inject
PublishSpec11ReportAction(
    @Config("projectId") String projectId,
    @Config("registryName") String registryName,
    @Parameter(ReportingModule.PARAM_JOB_ID) String jobId,
    Spec11EmailUtils emailUtils,
    Spec11RegistrarThreatMatchesParser spec11RegistrarThreatMatchesParser,
    Dataflow dataflow,
    Response response,
    @Parameter(PARAM_DATE) LocalDate date) {
  this.projectId = projectId;
  this.registryName = registryName;
  this.jobId = jobId;
  this.emailUtils = emailUtils;
  this.spec11RegistrarThreatMatchesParser = spec11RegistrarThreatMatchesParser;
  this.dataflow = dataflow;
  this.response = response;
  this.date = date;
}
 

@Inject
GenerateInvoicesAction(
    @Config("projectId") String projectId,
    @Config("apacheBeamBucketUrl") String beamBucketUrl,
    @Config("invoiceTemplateUrl") String invoiceTemplateUrl,
    @Config("defaultJobZone") String jobZone,
    @Parameter(PARAM_SHOULD_PUBLISH) boolean shouldPublish,
    YearMonth yearMonth,
    Dataflow dataflow,
    Response response,
    BillingEmailUtils emailUtils) {
  this.projectId = projectId;
  this.beamBucketUrl = beamBucketUrl;
  this.invoiceTemplateUrl = invoiceTemplateUrl;
  this.jobZone = jobZone;
  this.shouldPublish = shouldPublish;
  this.yearMonth = yearMonth;
  this.dataflow = dataflow;
  this.response = response;
  this.emailUtils = emailUtils;
}
 

@Before
public void setUp() throws IOException {
  response = new FakeResponse();
  dataflow = mock(Dataflow.class);

  // Establish the Dataflow API call chain
  dataflow = mock(Dataflow.class);
  dataflowProjects = mock(Dataflow.Projects.class);
  dataflowTemplates = mock(Templates.class);
  dataflowLaunch = mock(Launch.class);
  LaunchTemplateResponse launchTemplateResponse = new LaunchTemplateResponse();
  // Ultimately we get back this job response with a given id.
  launchTemplateResponse.setJob(new Job().setId("jobid"));
  when(dataflow.projects()).thenReturn(dataflowProjects);
  when(dataflowProjects.templates()).thenReturn(dataflowTemplates);
  when(dataflowTemplates.launch(any(String.class), any(LaunchTemplateParameters.class)))
      .thenReturn(dataflowLaunch);
  when(dataflowLaunch.setGcsPath(any(String.class))).thenReturn(dataflowLaunch);
  when(dataflowLaunch.execute()).thenReturn(launchTemplateResponse);
}
 

@Before
public void setUp() throws Exception {
  dataflow = mock(Dataflow.class);
  projects = mock(Projects.class);
  jobs = mock(Jobs.class);
  get = mock(Get.class);
  when(dataflow.projects()).thenReturn(projects);
  when(projects.jobs()).thenReturn(jobs);
  when(jobs.get("test-project", "12345")).thenReturn(get);
  expectedJob = new Job();
  when(get.execute()).thenReturn(expectedJob);
  emailUtils = mock(Spec11EmailUtils.class);
  parser = mock(Spec11RegistrarThreatMatchesParser.class);
  response = new FakeResponse();
  parser = mock(Spec11RegistrarThreatMatchesParser.class);
  publishAction =
      new PublishSpec11ReportAction(
          "test-project",
          "Super Cool Registry",
          "12345",
          emailUtils,
          parser,
          dataflow,
          response,
          date);
}
 

@Before
public void setUp() throws IOException {
  Dataflow dataflow = mock(Dataflow.class);
  Projects projects = mock(Projects.class);
  Jobs jobs = mock(Jobs.class);
  get = mock(Get.class);
  when(dataflow.projects()).thenReturn(projects);
  when(projects.jobs()).thenReturn(jobs);
  when(jobs.get("test-project", "12345")).thenReturn(get);
  expectedJob = new Job();
  when(get.execute()).thenReturn(expectedJob);
  emailUtils = mock(BillingEmailUtils.class);
  response = new FakeResponse();
  uploadAction =
      new PublishInvoicesAction(
          "test-project", "12345", emailUtils, dataflow, response, new YearMonth(2017, 10));
}
 

@Before
public void setUp() throws IOException {
  dataflow = mock(Dataflow.class);
  projects = mock(Projects.class);
  templates = mock(Templates.class);
  launch = mock(Launch.class);
  emailUtils = mock(BillingEmailUtils.class);
  when(dataflow.projects()).thenReturn(projects);
  when(projects.templates()).thenReturn(templates);
  when(templates.launch(any(String.class), any(LaunchTemplateParameters.class)))
      .thenReturn(launch);
  when(launch.setGcsPath(any(String.class))).thenReturn(launch);

  response = new FakeResponse();
  Job job = new Job();
  job.setId("12345");
  when(launch.execute()).thenReturn(new LaunchTemplateResponse().setJob(job));
}
 

/**
 * The root URL for the Dataflow API. {@code dataflowEndpoint} can override this value if it
 * contains an absolute URL, otherwise {@code apiRootUrl} will be combined with {@code
 * dataflowEndpoint} to generate the full URL to communicate with the Dataflow API.
 */
@Description(
    "The root URL for the Dataflow API. dataflowEndpoint can override this "
        + "value if it contains an absolute URL, otherwise apiRootUrl will be combined with "
        + "dataflowEndpoint to generate the full URL to communicate with the Dataflow API.")
@Default.String(Dataflow.DEFAULT_ROOT_URL)
String getApiRootUrl();
 

/**
 * An instance of the Dataflow client. Defaults to creating a Dataflow client using the current
 * set of options.
 */
@JsonIgnore
@Description(
    "An instance of the Dataflow client. Defaults to creating a Dataflow client "
        + "using the current set of options.")
@Default.InstanceFactory(DataflowClientFactory.class)
Dataflow getDataflowClient();
 

private static Dataflow buildMockDataflow(ArgumentMatcher<Job> jobMatcher) throws IOException {
  Dataflow mockDataflowClient = mock(Dataflow.class);
  Dataflow.Projects mockProjects = mock(Dataflow.Projects.class);
  Dataflow.Projects.Jobs mockJobs = mock(Dataflow.Projects.Jobs.class);
  Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class);

  when(mockDataflowClient.projects()).thenReturn(mockProjects);
  when(mockProjects.jobs()).thenReturn(mockJobs);
  when(mockJobs.create(eq("someProject"), argThat(jobMatcher))).thenReturn(mockRequest);

  Job resultJob = new Job();
  resultJob.setId("newid");
  when(mockRequest.execute()).thenReturn(resultJob);
  return mockDataflowClient;
}
 

@Test
public void testCancelTerminatedJob() throws IOException {
  Dataflow.Projects.Locations.Jobs.Get statusRequest =
      mock(Dataflow.Projects.Locations.Jobs.Get.class);

  Job statusResponse = new Job();
  statusResponse.setCurrentState("JOB_STATE_FAILED");
  when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest);
  when(statusRequest.execute()).thenReturn(statusResponse);

  Dataflow.Projects.Locations.Jobs.Update update =
      mock(Dataflow.Projects.Locations.Jobs.Update.class);
  when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class)))
      .thenReturn(update);
  when(update.execute()).thenThrow(new IOException());

  DataflowPipelineJob job =
      new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null);

  assertEquals(State.FAILED, job.cancel());
  Job content = new Job();
  content.setProjectId(PROJECT_ID);
  content.setId(JOB_ID);
  content.setRequestedState("JOB_STATE_CANCELLED");
  verify(mockJobs).update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), eq(content));
  verify(mockJobs).get(PROJECT_ID, REGION_ID, JOB_ID);
  verifyNoMoreInteractions(mockJobs);
}
 

private Dataflow buildMockDataflow() throws IOException {
  Dataflow mockDataflowClient = mock(Dataflow.class);
  Dataflow.Projects mockProjects = mock(Dataflow.Projects.class);
  Dataflow.Projects.Locations mockLocations = mock(Dataflow.Projects.Locations.class);
  Dataflow.Projects.Locations.Jobs.Create mockRequest =
      mock(Dataflow.Projects.Locations.Jobs.Create.class);
  Dataflow.Projects.Locations.Jobs.List mockList =
      mock(Dataflow.Projects.Locations.Jobs.List.class);

  when(mockDataflowClient.projects()).thenReturn(mockProjects);
  when(mockProjects.locations()).thenReturn(mockLocations);
  when(mockLocations.jobs()).thenReturn(mockJobs);
  when(mockJobs.create(eq(PROJECT_ID), eq(REGION_ID), isA(Job.class))).thenReturn(mockRequest);
  when(mockJobs.list(eq(PROJECT_ID), eq(REGION_ID))).thenReturn(mockList);
  when(mockList.setPageToken(any())).thenReturn(mockList);
  when(mockList.execute())
      .thenReturn(
          new ListJobsResponse()
              .setJobs(
                  Arrays.asList(
                      new Job()
                          .setName("oldjobname")
                          .setId("oldJobId")
                          .setCurrentState("JOB_STATE_RUNNING"))));

  Job resultJob = new Job();
  resultJob.setId("newid");
  when(mockRequest.execute()).thenReturn(resultJob);
  return mockDataflowClient;
}
 

@Test
public void testRunReturnDifferentRequestId() throws IOException {
  DataflowPipelineOptions options = buildPipelineOptions();
  Dataflow mockDataflowClient = options.getDataflowClient();
  Dataflow.Projects.Locations.Jobs.Create mockRequest =
      mock(Dataflow.Projects.Locations.Jobs.Create.class);
  when(mockDataflowClient
          .projects()
          .locations()
          .jobs()
          .create(eq(PROJECT_ID), eq(REGION_ID), any(Job.class)))
      .thenReturn(mockRequest);
  Job resultJob = new Job();
  resultJob.setId("newid");
  // Return a different request id.
  resultJob.setClientRequestId("different_request_id");
  when(mockRequest.execute()).thenReturn(resultJob);

  Pipeline p = buildDataflowPipeline(options);
  try {
    p.run();
    fail("Expected DataflowJobAlreadyExistsException");
  } catch (DataflowJobAlreadyExistsException expected) {
    assertThat(
        expected.getMessage(),
        containsString(
            "If you want to submit a second job, try again by setting a "
                + "different name using --jobName."));
    assertEquals(expected.getJob().getJobId(), resultJob.getId());
  }
}
 

@Test
public void testUpdateAlreadyUpdatedPipeline() throws IOException {
  DataflowPipelineOptions options = buildPipelineOptions();
  options.setUpdate(true);
  options.setJobName("oldJobName");
  Dataflow mockDataflowClient = options.getDataflowClient();
  Dataflow.Projects.Locations.Jobs.Create mockRequest =
      mock(Dataflow.Projects.Locations.Jobs.Create.class);
  when(mockDataflowClient
          .projects()
          .locations()
          .jobs()
          .create(eq(PROJECT_ID), eq(REGION_ID), any(Job.class)))
      .thenReturn(mockRequest);
  final Job resultJob = new Job();
  resultJob.setId("newid");
  // Return a different request id.
  resultJob.setClientRequestId("different_request_id");
  when(mockRequest.execute()).thenReturn(resultJob);

  Pipeline p = buildDataflowPipeline(options);

  thrown.expect(DataflowJobAlreadyUpdatedException.class);
  thrown.expect(
      new TypeSafeMatcher<DataflowJobAlreadyUpdatedException>() {
        @Override
        public void describeTo(Description description) {
          description.appendText("Expected job ID: " + resultJob.getId());
        }

        @Override
        protected boolean matchesSafely(DataflowJobAlreadyUpdatedException item) {
          return resultJob.getId().equals(item.getJob().getJobId());
        }
      });
  thrown.expectMessage(
      "The job named oldjobname with id: oldJobId has already been updated "
          + "into job id: newid and cannot be updated again.");
  p.run();
}
 

@Before
public void setUp() throws Exception {
  MockitoAnnotations.initMocks(this);
  when(transport.buildRequest(anyString(), anyString())).thenReturn(request);
  doCallRealMethod().when(request).getContentAsString();

  Dataflow service = new Dataflow(transport, Transport.getJsonFactory(), null);
  pipelineOptions = PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class);
  pipelineOptions.setProject(PROJECT_ID);
  pipelineOptions.setJobId(JOB_ID);
  pipelineOptions.setWorkerId(WORKER_ID);
  pipelineOptions.setGcpCredential(new TestCredential());
  pipelineOptions.setDataflowClient(service);
}
 

@Before
public void setUp() throws Exception {
  fakeGetConfigResponse = new GetDebugConfigResponse();
  fakeGetConfigResponse.setConfig(UPDATE_CONFIG_JSON);

  mockSendCapture = mock(Dataflow.Projects.Locations.Jobs.Debug.SendCapture.class);
  when(mockSendCapture.execute()).thenReturn(new SendDebugCaptureResponse());
}