def test_wait_gcp_resources_on_invalid_gcp_resource_uri(self, mock_build):
invalid_payload = '{"resources": [{"resourceType": "DataflowJob","resourceUri": "https://dataflow.googleapis.com/v1b3/projects/abc"}]}'
with self.assertRaises(ValueError):
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location,
invalid_payload,
self._gcp_resources_path)
def test_wait_gcp_resources_retries_to_get_status_on_non_completed_job(
self, mock_time_sleep, mock_build):
df_client = mock.Mock()
mock_build.return_value = df_client
expected_job_running = {
'id': 'job-1',
'currentState': 'JOB_STATE_RUNNING'
}
expected_job_succeeded = {
'id': 'job-1',
'currentState': 'JOB_STATE_DONE'
}
get_request = mock.Mock()
df_client.projects().locations().jobs().get.return_value = get_request
get_request.execute.side_effect = [
expected_job_running, expected_job_succeeded
]
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location, self._payload,
self._gcp_resources_path)
mock_time_sleep.assert_called_once_with(
wait_gcp_resources._POLLING_INTERVAL_IN_SECONDS)
self.assertEqual(df_client.projects().locations().jobs().get.call_count,
2)
def test_wait_gcp_resources_on_empty_gcp_resource(self, mock_build):
invalid_payload = '{"resources": [{}]}'
with self.assertRaises(ValueError):
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location,
invalid_payload,
self._gcp_resources_path)
def test_wait_gcp_resources_on_invalid_gcp_resource_type(self, mock_build):
invalid_payload = '{"resources": [{"resourceType": "BigQuery","resourceUri": "https://dataflow.googleapis.com/v1b3/projects/foo/locations/us-central1/jobs/job123"}]}'
with self.assertRaises(ValueError):
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location,
invalid_payload,
self._gcp_resources_path)
def test_wait_gcp_resources_cancel_dataflow(self,
mock_execution_context,
mock_build):
df_client = mock.Mock()
mock_build.return_value = df_client
expected_done_job = {'id': 'job-1', 'currentState': 'JOB_STATE_DONE'}
expected_running_job = {'id': 'job-1', 'currentState': 'JOB_STATE_RUNNING'}
get_request = mock.Mock()
df_client.projects().locations().jobs().get.return_value = get_request
# The first done_job is to end the polling loop.
get_request.execute.side_effect = [
expected_done_job, expected_running_job, expected_running_job
]
update_request = mock.Mock()
df_client.projects().locations().jobs().update.return_value = update_request
mock_execution_context.return_value = None
expected_cancel_job = {
'id': 'job-1',
'currentState': 'JOB_STATE_RUNNING',
'requestedState': 'JOB_STATE_CANCELLED'
}
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location, self._payload,
self._gcp_resources_path)
# Call cancellation handler
mock_execution_context.call_args[1]['on_cancel']()
df_client.projects().locations().jobs().update.assert_called_once_with(
projectId='foo', jobId='job123', location='us-central1',
body=expected_cancel_job)
def test_wait_gcp_resources_on_getting_failed_dataflow(self, mock_build):
df_client = mock.Mock()
mock_build.return_value = df_client
expected_job = {'id': 'job-1', 'currentState': 'JOB_STATE_FAILED'}
get_request = mock.Mock()
df_client.projects().locations().jobs().get.return_value = get_request
get_request.execute.return_value = expected_job
with self.assertRaises(RuntimeError):
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location, self._payload,
self._gcp_resources_path)
def test_wait_gcp_resources_on_getting_succeeded_dataflow(self, mock_build):
df_client = mock.Mock()
mock_build.return_value = df_client
expected_job = {'id': 'job-1', 'currentState': 'JOB_STATE_DONE'}
get_request = mock.Mock()
df_client.projects().locations().jobs().get.return_value = get_request
get_request.execute.return_value = expected_job
wait_gcp_resources.wait_gcp_resources(self._type, self._project,
self._location, self._payload,
self._gcp_resources_path)
df_client.projects().locations().jobs().get.assert_called_once_with(
projectId='foo', jobId='job123', location='us-central1', view=None)