def test_cancelled(self, mock_hook):
job = self.create_job(JobStatus.CANCELLED)
job_id = "job_id"
mock_hook.return_value.get_job.return_value = job
sensor = DataprocJobSensor(
task_id=TASK_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT,
dataproc_job_id=job_id,
gcp_conn_id=GCP_CONN_ID,
timeout=TIMEOUT,
)
with pytest.raises(AirflowException, match="Job was cancelled"):
sensor.poke(context={})
mock_hook.return_value.get_job.assert_called_once_with(
job_id=job_id, location=GCP_LOCATION, project_id=GCP_PROJECT)
def test_error(self, mock_hook):
job = self.create_job(JobStatus.ERROR)
job_id = "job_id"
mock_hook.return_value.get_job.return_value = job
sensor = DataprocJobSensor(
task_id=TASK_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT,
dataproc_job_id=job_id,
gcp_conn_id=GCP_CONN_ID,
timeout=TIMEOUT,
)
with self.assertRaisesRegex(AirflowException, "Job failed"):
sensor.poke(context={})
mock_hook.return_value.get_job.assert_called_once_with(
job_id=job_id, location=GCP_LOCATION, project_id=GCP_PROJECT)
def test_wait(self, mock_hook):
job = self.create_job(JobStatus.RUNNING)
job_id = "job_id"
mock_hook.return_value.get_job.return_value = job
sensor = DataprocJobSensor(
task_id=TASK_ID,
location=GCP_LOCATION,
project_id=GCP_PROJECT,
dataproc_job_id=job_id,
gcp_conn_id=GCP_CONN_ID,
timeout=TIMEOUT,
)
ret = sensor.poke(context={})
mock_hook.return_value.get_job.assert_called_once_with(
job_id=job_id, location=GCP_LOCATION, project_id=GCP_PROJECT)
assert not ret