def test_dataflow_job_init_without_job_id(self): mock_jobs = MagicMock() self.mock_dataflow.projects.return_value.locations.return_value. \ jobs.return_value = mock_jobs _DataflowJob(self.mock_dataflow, TEST_PROJECT, TEST_JOB_NAME, TEST_LOCATION, 10) mock_jobs.list.assert_called_once_with(projectId=TEST_PROJECT, location=TEST_LOCATION)
def test_dataflow_job_wait_for_single_jobs(self): job = { "id": TEST_JOB_ID, "name": TEST_JOB_NAME, "currentState": DataflowJobStatus.JOB_STATE_DONE } self.mock_dataflow.projects.return_value.locations.return_value. \ jobs.return_value.get.return_value.execute.return_value = job dataflow_job = _DataflowJob(dataflow=self.mock_dataflow, project_number=TEST_PROJECT, name=TEST_JOB_NAME, location=TEST_LOCATION, poll_sleep=10, job_id=TEST_JOB_ID, num_retries=20, multiple_jobs=False) dataflow_job.wait_for_done() self.mock_dataflow.projects.return_value.locations.return_value. \ jobs.return_value.get.assert_called_once_with( jobId=TEST_JOB_ID, location=TEST_LOCATION, projectId=TEST_PROJECT ) self.mock_dataflow.projects.return_value.locations.return_value. \ jobs.return_value.get.return_value.execute.assert_called_once_with(num_retries=20) self.assertEqual(dataflow_job.get(), [job])
def test_dataflow_job_wait_for_multiple_jobs_and_one_unknown(self): (self.mock_dataflow.projects.return_value.locations.return_value.jobs. return_value.list.return_value.execute.return_value) = { "jobs": [{ "id": "id-1", "name": "name-1", "currentState": DataflowJobStatus.JOB_STATE_DONE }, { "id": "id-2", "name": "name-2", "currentState": "unknown" }] } dataflow_job = _DataflowJob(dataflow=self.mock_dataflow, project_number=TEST_PROJECT, name="name-", location=TEST_LOCATION, poll_sleep=0, job_id=None, num_retries=20, multiple_jobs=True) with self.assertRaisesRegex( Exception, 'Google Cloud Dataflow job name-2 was unknown state: unknown'): dataflow_job.wait_for_done()
def test_dataflow_job_wait_for_multiple_jobs_and_streaming_jobs(self): mock_jobs_list = (self.mock_dataflow.projects.return_value.locations. return_value.jobs.return_value.list) mock_jobs_list.return_value.execute.return_value = { "jobs": [{ "id": "id-2", "name": "name-2", "currentState": DataflowJobStatus.JOB_STATE_RUNNING, "type": DataflowJobType.JOB_TYPE_STREAMING }] } dataflow_job = _DataflowJob(dataflow=self.mock_dataflow, project_number=TEST_PROJECT, name="name-", location=TEST_LOCATION, poll_sleep=0, job_id=None, num_retries=20, multiple_jobs=True) dataflow_job.wait_for_done() self.assertEqual(1, mock_jobs_list.call_count)