def test_launch_python_no_job_created(self, mock_subprocess, mock_process,
mock_client, mock_context,
mock_stage_file, mock_display,
mock_storage):
mock_context().__enter__().context_id.return_value = 'ctx-1'
mock_process().read_lines.return_value = [b'no job id', b'no job id']
result = launch_python('/tmp/test.py', 'project-1', 'us-central1')
self.assertEqual(None, result)
def test_launch_python_succeed(self, mock_subprocess, mock_process,
mock_client, mock_context, mock_stage_file,
mock_display):
mock_context().__enter__().context_id.return_value = 'ctx-1'
mock_client().list_aggregated_jobs.return_value = {'jobs': []}
mock_process().read_lines.return_value = [
b'https://console.cloud.google.com/dataflow/locations/us-central1/jobs/job-1?project=project-1'
]
expected_job = {'id': 'job-1', 'currentState': 'JOB_STATE_DONE'}
mock_client().get_job.return_value = expected_job
result = launch_python('/tmp/test.py', 'project-1')
self.assertEqual(expected_job, result)
def test_launch_python_retry_succeed(self, mock_subprocess, mock_process,
mock_client, mock_context, mock_stage_file, mock_display, mock_storage):
mock_context().__enter__().context_id.return_value = 'ctx-1'
mock_storage.Client().bucket().blob().exists.return_value = True
mock_storage.Client().bucket().blob().download_as_string.return_value = b'job-1,us-central1'
expected_job = {
'id': 'job-1',
'currentState': 'JOB_STATE_DONE'
}
mock_client().get_job.return_value = expected_job
result = launch_python('/tmp/test.py', 'project-1', 'us-central1', staging_dir='gs://staging/dir')
self.assertEqual(expected_job, result)
mock_process.assert_not_called()
def test_launch_python_retry_succeed(self, mock_subprocess, mock_process,
mock_client, mock_context, mock_stage_file, mock_display):
mock_context().__enter__().context_id.return_value = 'ctx-1'
mock_client().list_aggregated_jobs.return_value = {
'jobs': [{
'id': 'job-1',
'name': 'test_job-ctx-1'
}]
}
expected_job = {
'currentState': 'JOB_STATE_DONE'
}
mock_client().get_job.return_value = expected_job
result = launch_python('/tmp/test.py', 'project-1', job_name_prefix='test-job')
self.assertEqual(expected_job, result)
mock_process.assert_not_called()
def test_launch_python_succeed(self, mock_subprocess, mock_process,
mock_client, mock_context, mock_stage_file,
mock_display, mock_storage):
mock_context().__enter__().context_id.return_value = 'ctx-1'
mock_storage.Client().bucket().blob().exists.return_value = False
mock_process().read_lines.return_value = [
b'https://console.cloud.google.com/dataflow/locations/us-central1/jobs/job-1?project=project-1'
]
expected_job = {'id': 'job-1', 'currentState': 'JOB_STATE_DONE'}
mock_client().get_job.return_value = expected_job
result = launch_python('/tmp/test.py',
'project-1',
staging_dir='gs://staging/dir')
self.assertEqual(expected_job, result)
mock_storage.Client().bucket().blob(
).upload_from_string.assert_called_with('job-1,us-central1')
