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
)
def test_dataflow_job_init_with_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, TEST_JOB_ID).get_jobs()
mock_jobs.get.assert_called_once_with(projectId=TEST_PROJECT, location=TEST_LOCATION,
jobId=TEST_JOB_ID)
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_wait_for_multiple_jobs(self):
job = {
"id": TEST_JOB_ID,
"name": UNIQUE_JOB_NAME,
"currentState": DataflowJobStatus.JOB_STATE_DONE
}
# fmt: off
(self.mock_dataflow.projects.return_value.locations.return_value.jobs.
return_value.list.return_value.execute.return_value) = {
"jobs": [job, job]
}
(self.mock_dataflow.projects.return_value.locations.return_value.jobs.
return_value.list_next.return_value) = None
# fmt: on
dataflow_job = _DataflowJobsController(
dataflow=self.mock_dataflow,
project_number=TEST_PROJECT,
name=UNIQUE_JOB_NAME,
location=TEST_LOCATION,
poll_sleep=10,
job_id=TEST_JOB_ID,
num_retries=20,
multiple_jobs=True,
)
dataflow_job.wait_for_done()
# fmt: off
self.mock_dataflow.projects.return_value.locations.return_value.jobs.return_value.\
list.assert_called_once_with(location=TEST_LOCATION, projectId=TEST_PROJECT)
self.mock_dataflow.projects.return_value.locations.return_value.jobs.return_value.list\
.return_value.execute.assert_called_once_with(num_retries=20)
# fmt: on
self.assertEqual(dataflow_job.get_jobs(), [job, job])
def test_dataflow_job_is_job_running_with_no_job(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": []
}
(
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
)
result = dataflow_job.is_job_running()
self.assertEqual(False, result)
def test_dataflow_job_cancel_job(self):
job = {
"id": TEST_JOB_ID, "name": UNIQUE_JOB_NAME, "currentState": DataflowJobStatus.JOB_STATE_RUNNING
}
get_method = (
self.mock_dataflow.projects.return_value.
locations.return_value.
jobs.return_value.
get
)
get_method.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=UNIQUE_JOB_NAME,
location=TEST_LOCATION,
poll_sleep=10,
job_id=TEST_JOB_ID,
num_retries=20,
multiple_jobs=False
)
dataflow_job.cancel()
get_method.assert_called_once_with(
jobId=TEST_JOB_ID,
location=TEST_LOCATION,
projectId=TEST_PROJECT
)
get_method.return_value.execute.assert_called_once_with(num_retries=20)
self.mock_dataflow.new_batch_http_request.assert_called_once_with()
mock_batch = self.mock_dataflow.new_batch_http_request.return_value
mock_update = (
self.mock_dataflow.projects.return_value.
locations.return_value.
jobs.return_value.
update
)
mock_update.assert_called_once_with(
body={'requestedState': 'JOB_STATE_CANCELLED'},
jobId='test-job-id',
location=TEST_LOCATION,
projectId='test-project',
)
mock_batch.add.assert_called_once_with(
mock_update.return_value
)
mock_batch.execute.assert_called_once()
