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
_DataflowJobsController(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
(self.mock_dataflow.projects.return_value.locations.return_value.jobs.
return_value.list_next.return_value) = None
dataflow_job = _DataflowJobsController(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_jobs(), [job])
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
}]
}
(self.mock_dataflow.projects.return_value.locations.return_value.jobs.
return_value.list_next.return_value) = None
dataflow_job = _DataflowJobsController(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)
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"
}]
}
(self.mock_dataflow.projects.return_value.locations.return_value.jobs.
return_value.list_next.return_value) = None
dataflow_job = _DataflowJobsController(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_init_without_job_id(self):
job = {
"id": TEST_JOB_ID,
"name": TEST_JOB_NAME,
"currentState": DataflowJobStatus.JOB_STATE_DONE
}
mock_list = (self.mock_dataflow.projects.return_value.locations.
return_value.jobs.return_value.list)
(mock_list.return_value.execute.return_value) = {'jobs': [job]}
(self.mock_dataflow.projects.return_value.locations.return_value.jobs.
return_value.list_next.return_value) = None
_DataflowJobsController(self.mock_dataflow, TEST_PROJECT,
TEST_JOB_NAME, TEST_LOCATION, 10).get_jobs()
mock_list.assert_called_once_with(projectId=TEST_PROJECT,
location=TEST_LOCATION)