def cancel_job(beam_job_run_model: beam_job_models.BeamJobRunModel) -> None:
"""Refreshs the state of the given BeamJobRunModel.
Args:
beam_job_run_model: BeamJobRunModel. The model to update.
Raises:
ValueError. The given model has no job ID.
"""
job_id = beam_job_run_model.dataflow_job_id
if job_id is None:
raise ValueError('dataflow_job_id must not be None')
try:
dataflow.JobsV1Beta3Client().update_job(
dataflow.UpdateJobRequest(
job_id=job_id,
project_id=feconf.OPPIA_PROJECT_ID,
location=feconf.GOOGLE_APP_ENGINE_REGION,
job=dataflow.Job(
requested_state=dataflow.JobState.JOB_STATE_CANCELLED)))
except Exception:
logging.exception('Failed to cancel job_id="%s"!' % job_id)
else:
beam_job_run_model.latest_job_state = (
beam_job_models.BeamJobState.CANCELLING.value)
beam_job_run_model.update_timestamps()
def setUp(self) -> None:
super().setUp()
self.run_model = beam_job_services.create_beam_job_run_model(
'WorkingJob', dataflow_job_id='123')
self.dataflow_job = dataflow.Job(
id='123',
project_id=feconf.OPPIA_PROJECT_ID,
location=feconf.GOOGLE_APP_ENGINE_REGION,
current_state=dataflow.JobState.JOB_STATE_CANCELLING,
current_state_time=datetime.datetime.utcnow())
self.dataflow_client_mock = mock.Mock()
self.dataflow_client_mock.update_job.return_value = self.dataflow_job
self.exit_stack = contextlib.ExitStack()
self.exit_stack.enter_context(self.swap_to_always_return(
dataflow, 'JobsV1Beta3Client', value=self.dataflow_client_mock))