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