amoghrajesh commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4492991937
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -7422,39 +7423,135 @@ def execute(self, context):
]
assert ClearTaskStateStore not in sent_types
+ @staticmethod
+ @contextlib.contextmanager
+ def _patch_state_backend(backend):
+ # Both modules resolve the backend independently, so both lookups need
patching.
+ with (
+ mock.patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=backend,
+ ),
+ mock.patch(
+
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
+ return_value=backend,
+ ),
+ ):
+ yield backend
+
+ @staticmethod
+ def _run_and_finalize(runtime_ti, log):
+ # The clear lives at the tail of finalize(), so run() alone would not
reach it.
+ context = runtime_ti.get_template_context()
+ state, msg, error = run(runtime_ti, context=context, log=log)
+ finalize(runtime_ti, state, context, log, error)
+ return state, msg, error
+
@conf_vars({("state_store", "clear_on_success"): "True"})
def test_clear_on_success_clears_backend_without_comms_roundtrip(
self, create_runtime_ti, mock_supervisor_comms
):
"""clear_on_success calls backend.clear() directly without sending
ClearTaskStateStore comms."""
- mock_backend = mock.MagicMock()
+ backend = mock.create_autospec(BaseStoreBackend, instance=True)
class MyOperator(BaseOperator):
def execute(self, context):
pass
- task = MyOperator(task_id="t")
- runtime_ti = create_runtime_ti(task=task)
+ runtime_ti = create_runtime_ti(task=MyOperator(task_id="t"))
- with (
- mock.patch(
-
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
- return_value=mock_backend,
- ),
- mock.patch(
-
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
- return_value=mock_backend,
- ),
- ):
- run(runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock())
+ with self._patch_state_backend(backend):
+ self._run_and_finalize(runtime_ti, mock.MagicMock())
- mock_backend.clear.assert_called_once()
+ backend.clear.assert_called_once()
sent_types = [
type(call.kwargs.get("msg") or (call.args[0] if call.args else
None))
for call in mock_supervisor_comms.send.call_args_list
]
assert ClearTaskStateStore not in sent_types
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_clear_on_success_runs_after_success_is_reported(self,
create_runtime_ti, mock_supervisor_comms):
+ """The backend is only cleared once SucceedTask has been reported to
the supervisor."""
+ backend = mock.create_autospec(BaseStoreBackend, instance=True)
+ order = []
+ mock_supervisor_comms.send.side_effect = lambda msg, **kwargs:
order.append(type(msg).__name__)
+ backend.clear.side_effect = lambda *a, **kw: order.append("clear")
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ pass
+
+ runtime_ti = create_runtime_ti(task=MyOperator(task_id="t"))
+
+ with self._patch_state_backend(backend):
+ state, _, _ = self._run_and_finalize(runtime_ti, mock.MagicMock())
+
+ assert state == TaskInstanceState.SUCCESS
+ assert order.index("SucceedTask") < order.index("clear")
Review Comment:
Added an `on_success_callback` that appends to order; the test now asserts
`SucceedTask` < `callback` < `clear`.
--
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]
