def test_refresh_state_of_all_beam_job_run_models_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_services.refresh_state_of_all_beam_job_run_models()
# 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.assertEqual(
updated_beam_job_run_models,
beam_job_models.BeamJobRunModel.get_multi(
[m.id for m in initial_beam_job_run_models]))
def test_refresh_state_of_all_beam_job_run_models(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_services.refresh_state_of_all_beam_job_run_models()
# 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.assertEqual(
updated_beam_job_run_models,
beam_job_models.BeamJobRunModel.get_multi(
[m.id for m in initial_beam_job_run_models]))