abhijeets25012-tech commented on code in PR #76184:
URL: https://github.com/apache/airflow/pull/76184#discussion_r4500971314
##########
task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py:
##########
@@ -730,6 +730,56 @@ def test_start_imports_task_runner_before_forking(self,
base_start_kwargs):
assert calls == ["import task runner", "fork"]
+ def test_team_name_is_set_before_the_callback_runs(self,
base_start_kwargs):
+ """Secret lookups in callback code must resolve to the owning team,
even at import time."""
+ calls = []
+ self.mock_execute_callback.side_effect = lambda **kwargs:
calls.append("execute") or (True, None)
+
+ with patch(
+
"airflow.sdk.execution_time.callback_supervisor.set_current_team_name",
+ side_effect=lambda team_name: calls.append(("team", team_name)),
+ ):
+ CallbackSubprocess.start(**{**base_start_kwargs, "team_name":
"team-a"})
+ self.mock_super_start.call_args.kwargs["target"]()
+
+ assert calls == [("team", "team-a"), "execute"]
Review Comment:
Good catch — applied. The test now records bundle initialisation too and
passes a `bundle_info`, so it pins the intended order:
```python
def test_team_name_is_set_before_the_callback_runs(self, base_start_kwargs,
mock_bundle_setup):
calls = []
mock_bundle_setup["bundle"].initialize.side_effect = lambda:
calls.append("bundle init")
...
assert calls == [("team", "team-a"), "bundle init", "execute"]
```
Verified it does what you asked: moving `set_current_team_name(team_name)`
below the bundle block now fails this test (the call list becomes `["bundle
init", ("team", "team-a"), "execute"]`), and it passes with the current order.
43 passed in the file.
Generated-by: Qoder Agent (Claude) following the guidelines.
##########
task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py:
##########
@@ -730,6 +730,56 @@ def test_start_imports_task_runner_before_forking(self,
base_start_kwargs):
assert calls == ["import task runner", "fork"]
+ def test_team_name_is_set_before_the_callback_runs(self,
base_start_kwargs):
+ """Secret lookups in callback code must resolve to the owning team,
even at import time."""
+ calls = []
+ self.mock_execute_callback.side_effect = lambda **kwargs:
calls.append("execute") or (True, None)
+
+ with patch(
+
"airflow.sdk.execution_time.callback_supervisor.set_current_team_name",
+ side_effect=lambda team_name: calls.append(("team", team_name)),
+ ):
+ CallbackSubprocess.start(**{**base_start_kwargs, "team_name":
"team-a"})
+ self.mock_super_start.call_args.kwargs["target"]()
+
+ assert calls == [("team", "team-a"), "execute"]
+
+
+class TestCallbackTeamSecrets:
+ """A callback subprocess must resolve secrets for the Dag run's team."""
+
+ def test_team_variable_wins_over_the_global_one(self, monkeypatch):
Review Comment:
Applied — the test now takes `mock_supervisor_comms`, so the
`task_runner.SUPERVISOR_COMMS` set by the real `_target` is restored by the
fixture afterwards instead of leaking a MagicMock into later tests.
Generated-by: Qoder Agent (Claude) following the guidelines.
--
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]
