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()