def test_get_beam_job_runs_with_refresh_when_command_fails(self):
initial_beam_job_run_models = [
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.DONE.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.RUNNING.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.CANCELLED.value),
]
beam_job_models.BeamJobRunModel.update_timestamps_multi(
initial_beam_job_run_models)
beam_job_models.BeamJobRunModel.put_multi(initial_beam_job_run_models)
gcloud_failure_mock = self.swap_to_always_raise(
subprocess,
'check_output',
error=subprocess.CalledProcessError(1, 'gcloud', 'Uh-oh!'))
with gcloud_failure_mock, self.capture_logging() as logs:
self.assertRaisesRegexp(
subprocess.CalledProcessError, 'returned non-zero exit status',
lambda: beam_job_services.get_beam_job_runs(refresh=True))
self.assertEqual(
initial_beam_job_run_models,
beam_job_models.BeamJobRunModel.get_multi(
[m.id for m in initial_beam_job_run_models]))
self.assertIn('Failed to update the state of job', logs[-1])
def test_get_beam_job_runs_with_refresh_of_synchronous_job(self):
initial_beam_job_run_models = [
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.DONE.value),
self.create_beam_job_run_model(
# Synchronous jobs do not create Dataflow Jobs, so the value of
# the dataflow_job_id is None.
dataflow_job_id=None,
job_state=beam_job_models.BeamJobState.RUNNING.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.CANCELLED.value),
]
beam_job_models.BeamJobRunModel.update_timestamps_multi(
initial_beam_job_run_models)
beam_job_models.BeamJobRunModel.put_multi(initial_beam_job_run_models)
utcnow = datetime.datetime.utcnow()
with self.mock_datetime_utcnow(utcnow):
beam_job_runs = beam_job_services.get_beam_job_runs(refresh=True)
# Only the second model (job_state=RUNNING) should have been updated,
# the other two (DONE and CANCELLED) are in terminal states and don't
# need to be checked.
updated_beam_job_run_models = initial_beam_job_run_models[:]
updated_beam_job_run_models[1].latest_job_state = (
beam_job_models.BeamJobState.UNKNOWN.value)
updated_beam_job_run_models[1].last_updated = utcnow
self.assert_domains_equal_models(beam_job_runs,
updated_beam_job_run_models)
def test_get_beam_job_runs_with_refresh(self):
initial_beam_job_run_models = [
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.DONE.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.RUNNING.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.CANCELLED.value),
]
beam_job_models.BeamJobRunModel.update_timestamps_multi(
initial_beam_job_run_models)
beam_job_models.BeamJobRunModel.put_multi(initial_beam_job_run_models)
gcloud_output_mock = self.swap_check_output_with_gcloud_response(
current_state='JOB_STATE_DONE',
current_state_time=self.NANO_RESOLUTION_ISO_8601_STR)
with gcloud_output_mock:
beam_job_runs = beam_job_services.get_beam_job_runs(refresh=True)
# Only the second model (job_state=RUNNING) should have been updated,
# the other two (DONE and CANCELLED) are in terminal states and don't
# need to be checked.
updated_beam_job_run_models = initial_beam_job_run_models[:]
updated_beam_job_run_models[1].latest_job_state = (
beam_job_models.BeamJobState.DONE.value)
updated_beam_job_run_models[1].last_updated = (datetime.datetime(
2014, 10, 2, 15, 1, 23, 45123))
self.assert_domains_equal_models(beam_job_runs,
updated_beam_job_run_models)
def test_create_beam_job_run_model(self) -> None:
model = beam_job_services.create_beam_job_run_model(
'FooJob', dataflow_job_id='123')
model.put()
all_runs = beam_job_services.get_beam_job_runs(refresh=False)
self.assertEqual(len(all_runs), 1)
run = all_runs[0]
self.assertEqual(run.job_name, 'FooJob')
self.assertFalse(run.job_is_synchronous)
def test_create_beam_job_run_model(self):
model = beam_job_services.create_beam_job_run_model(
'FooJob', ['num_foos'], dataflow_job_id='123')
model.put()
all_runs = beam_job_services.get_beam_job_runs()
self.assertEqual(len(all_runs), 1)
run = all_runs[0]
self.assertEqual(run.job_name, 'FooJob')
self.assertEqual(run.job_arguments, ['num_foos'])
self.assertFalse(run.job_is_synchronous)
def test_get_beam_job_runs(self):
beam_job_run_models = [
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.DONE.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.RUNNING.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.CANCELLED.value),
]
beam_job_models.BeamJobRunModel.update_timestamps_multi(
beam_job_run_models)
beam_job_models.BeamJobRunModel.put_multi(beam_job_run_models)
self.assert_domains_equal_models(beam_job_services.get_beam_job_runs(),
beam_job_run_models)
def test_get_beam_job_runs_with_refresh(self) -> None:
beam_job_run_models = [
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.DONE.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.RUNNING.value),
self.create_beam_job_run_model(
job_state=beam_job_models.BeamJobState.CANCELLED.value),
]
beam_job_models.BeamJobRunModel.update_timestamps_multi(
beam_job_run_models)
beam_job_models.BeamJobRunModel.put_multi(beam_job_run_models)
with self.swap_to_always_return(jobs_manager,
'refresh_state_of_beam_job_run_model'):
self.assert_domains_equal_models(
beam_job_services.get_beam_job_runs(refresh=True),
beam_job_run_models)
def get(self) -> None:
sorted_beam_job_runs = sorted(beam_job_services.get_beam_job_runs(),
key=lambda j: j.job_updated_on,
reverse=True)
self.render_json({'runs': [r.to_dict() for r in sorted_beam_job_runs]})
