amoghrajesh commented on code in PR #71211:
URL: https://github.com/apache/airflow/pull/71211#discussion_r3756729227
##########
providers/amazon/tests/unit/amazon/aws/operators/test_glue.py:
##########
@@ -781,6 +784,295 @@ def test_inject_parent_job_info_with_resume_on_retry(
assert GlueJobOperator.TASK_UUID_ARG in call_args
+class TestGlueJobOperatorDeprecation:
+ @pytest.mark.parametrize("resume_value", [True, False])
+ def test_warns_on_3_3_plus_and_maps_to_durable(self, resume_value):
+ with
mock.patch("airflow.providers.amazon.aws.operators.glue.AIRFLOW_V_3_3_PLUS",
True):
+ with pytest.warns(AirflowProviderDeprecationWarning,
match="resume_glue_job_on_retry"):
+ glue = GlueJobOperator(
+ task_id=TASK_ID, job_name=JOB_NAME,
resume_glue_job_on_retry=resume_value
+ )
+ assert glue.durable is resume_value
+
+ def test_silent_below_3_3(self):
+ with
mock.patch("airflow.providers.amazon.aws.operators.glue.AIRFLOW_V_3_3_PLUS",
False):
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
+ glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME,
resume_glue_job_on_retry=True)
+ deprecation_warnings = [
+ w for w in caught if issubclass(w.category,
AirflowProviderDeprecationWarning)
+ ]
+ assert deprecation_warnings == []
+ assert glue.durable is True
+
+ def test_both_flags_passed_durable_wins(self):
+ with
mock.patch("airflow.providers.amazon.aws.operators.glue.AIRFLOW_V_3_3_PLUS",
True):
+ with pytest.warns(AirflowProviderDeprecationWarning):
+ glue = GlueJobOperator(
+ task_id=TASK_ID,
+ job_name=JOB_NAME,
+ durable=True,
+ resume_glue_job_on_retry=False,
+ )
+ assert glue.durable is True
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS,
+ reason="The <3.3 compat stub's __init__ isn't decorated with
BaseOperatorMeta._apply_defaults, "
+ "so default_args injection for durable only works on the real
ResumableJobMixin.",
+ )
+ def test_default_args_durable_reaches_operator(self):
+ with DAG(
+ dag_id="test_glue_durable_default_args",
+ schedule=None,
+ start_date=datetime(2024, 1, 1),
+ default_args={"durable": False},
+ ):
+ glue = GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME)
+ assert glue.durable is False
+
+
+class FakeTaskStateStore:
+ """In-memory task state store for tests."""
+
+ def __init__(self, stored: dict[str, str] | None = None):
+ self._store: dict[str, str] = dict(stored or {})
+
+ def get(self, key: str) -> str | None:
+ return self._store.get(key)
+
+ def set(self, key: str, value: str) -> None:
+ self._store[key] = value
+
+
[email protected](
+ not AIRFLOW_V_3_3_PLUS,
+ reason="ResumableJobMixin reconnect requires task_state_store, available
in Airflow 3.3+",
+)
+class TestGlueJobOperatorDurableExecution:
+ def _build(self, **kwargs):
+ return GlueJobOperator(task_id=TASK_ID, job_name=JOB_NAME, **kwargs)
+
+ def _stub_empty_scan(self, glue):
+ # submit_job scans for a task UUID tagged run whenever durable is set,
regardless of why it
+ # was called. Stub it to return no matches so tests don't depend on
that fallback mechanism.
+ glue.hook.conn = mock.MagicMock()
+ glue.hook.conn.get_job_runs.return_value = {"JobRuns": []}
+
+ def _context(self, store=None, try_number=2):
+ ti = mock.MagicMock()
+ ti.try_number = try_number
+ ti.xcom_pull.return_value = None
+ ctx = {"ti": ti}
+ if store is not None:
+ ctx["task_state_store"] = store
+ return ctx
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_fresh_submit_persists_before_polling(
+ self, mock_get_conn, mock_initialize_job, mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ self._stub_empty_scan(glue)
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ store = FakeTaskStateStore()
+ persisted_before_poll = []
+ mock_job_completion.side_effect = lambda *a, **k: (
+ persisted_before_poll.append(store.get("glue_job_run_id")) or
{"JobRunState": "SUCCEEDED"}
+ )
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ assert store.get("glue_job_run_id") == "jr_new"
+ assert persisted_before_poll == ["jr_new"]
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_reconnect_when_stored_run_is_running(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ mock_get_job_state.return_value = "RUNNING"
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_old"
+ mock_initialize_job.assert_not_called()
+ mock_job_completion.assert_called_once_with(JOB_NAME, "jr_old", False,
0)
+
+ @pytest.mark.parametrize("status", ["STARTING", "RUNNING", "WAITING",
"STOPPING"])
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_reconnect_from_every_active_state(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion, status
+ ):
+ glue = self._build(durable=True)
+ mock_get_job_state.return_value = status
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ glue.execute(self._context(store))
+
+ mock_initialize_job.assert_not_called()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_already_succeeded_returns_without_resubmit(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ mock_get_job_state.return_value = "SUCCEEDED"
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_old"
+ mock_initialize_job.assert_not_called()
+ mock_job_completion.assert_not_called()
+
+ @pytest.mark.parametrize("status", ["FAILED", "TIMEOUT", "STOPPED"])
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_terminal_failure_resubmits_fresh(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion, status
+ ):
+ """STOPPED resubmits like any other terminal state: Glue's API can't
tell a console
+ cancellation from a run this operator's own on_kill stopped, so
treating it as success
+ would risk silently reporting a self-inflicted stop as done."""
+ glue = self._build(durable=True)
+ self._stub_empty_scan(glue)
+ mock_get_job_state.return_value = status
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ assert store.get("glue_job_run_id") == "jr_new"
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "get_job_state")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_not_found_resubmits_fresh(
+ self, mock_get_conn, mock_initialize_job, mock_get_job_state,
mock_job_completion
+ ):
+ glue = self._build(durable=True)
+ self._stub_empty_scan(glue)
+ mock_get_job_state.side_effect = ClientError(
+ {"Error": {"Code": "EntityNotFoundException", "Message": "gone"}},
"GetJobRun"
+ )
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_durable_false_never_touches_store(self, mock_get_conn,
mock_initialize_job, mock_job_completion):
+ glue = self._build(durable=False)
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore({"glue_job_run_id": "jr_old"})
+
+ job_run_id = glue.execute(self._context(store))
+
+ assert job_run_id == "jr_new"
+ assert store.get("glue_job_run_id") == "jr_old", "store must be left
untouched"
+ mock_initialize_job.assert_called_once()
+
+ @mock.patch.object(GlueJobHook, "conn", new_callable=mock.PropertyMock)
+ @mock.patch.object(GlueJobHook, "job_completion")
+ @mock.patch.object(GlueJobHook, "initialize_job")
+ @mock.patch.object(GlueJobHook, "get_conn")
+ def test_first_attempt_skips_the_retry_lookup_entirely(
+ self, mock_get_conn, mock_initialize_job, mock_job_completion,
mock_conn
+ ):
+ glue = self._build(durable=True)
+ mock_initialize_job.return_value = {"JobRunId": "jr_new"}
+ mock_job_completion.return_value = {"JobRunState": "SUCCEEDED"}
+ store = FakeTaskStateStore()
+
+ job_run_id = glue.execute(self._context(store, try_number=1))
+
+ assert job_run_id == "jr_new"
+ mock_initialize_job.assert_called_once()
+ mock_conn.return_value.get_job_run.assert_not_called()
+ mock_conn.return_value.get_job_runs.assert_not_called()
Review Comment:
Added.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
