seanghaeli commented on code in PR #76184:
URL: https://github.com/apache/airflow/pull/76184#discussion_r4496918819
##########
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:
Found this in testing: the comment in `_target` says the team is set before
the bundle init so a team's git connection gets used, but moving that line
below the bundle block still passes all 43 tests. This would pin the order. Not
blocking.
```suggestion
def test_team_name_is_set_before_the_callback_runs(self,
base_start_kwargs, mock_bundle_setup):
"""Secret lookups in callback code must resolve to the owning team,
even at import time."""
calls = []
mock_bundle_setup["bundle"].initialize.side_effect = lambda:
calls.append("bundle init")
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)),
):
bundle_info = BundleInfo(name="test-bundle", version="1.0")
CallbackSubprocess.start(**{**base_start_kwargs, "team_name":
"team-a", "bundle_info": bundle_info})
self.mock_super_start.call_args.kwargs["target"]()
assert calls == [("team", "team-a"), "bundle init", "execute"]
```
##########
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:
This runs the real `_target`, which sets `task_runner.SUPERVISOR_COMMS`, but
nothing resets it afterwards, so it stays a MagicMock for later tests. The
`mock_supervisor_comms` fixture should clean that up. Not blocking.
```suggestion
def test_team_variable_wins_over_the_global_one(self, monkeypatch,
mock_supervisor_comms):
```
--
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]
