rjgoyln commented on code in PR #73473:
URL: https://github.com/apache/airflow/pull/73473#discussion_r4115545837
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -6827,3 +6827,28 @@ def
test_stats_tags_with_standalone_and_key_value_tags(create_runtime_ti):
"task_id": "t",
"run_type": "manual",
}
+
+
+def test_post_execute_receives_result(create_runtime_ti,
mock_supervisor_comms):
+ """Verify that an overridden post_execute receives the result of the
execute method."""
+ post_execute_results: list[str] = []
+
+ from airflow.sdk import BaseOperator
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ return "my_success_result"
+
+ def post_execute(self, context, result=None):
+ post_execute_results.append(result)
+
+ task = MyOperator(task_id="post_execute_task")
+ ti = create_runtime_ti(task=task)
+ import unittest.mock as mock
+
+ log = mock.MagicMock(spec=["info", "debug", "warning", "error",
"exception", "bind"])
+
+ run(ti, context=ti.get_template_context(), log=log)
+
+ assert post_execute_results == ["my_success_result"]
+
Review Comment:
The file gains a second trailing newline here, which `ruff-format` strips —
`ruff format --diff` on this file reports exactly this line, so the
static-checks job fails until it is dropped.
```suggestion
assert post_execute_results == ["my_success_result"]
```
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -6827,3 +6827,28 @@ def
test_stats_tags_with_standalone_and_key_value_tags(create_runtime_ti):
"task_id": "t",
"run_type": "manual",
}
+
+
+def test_post_execute_receives_result(create_runtime_ti,
mock_supervisor_comms):
+ """Verify that an overridden post_execute receives the result of the
execute method."""
+ post_execute_results: list[str] = []
+
+ from airflow.sdk import BaseOperator
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ return "my_success_result"
+
+ def post_execute(self, context, result=None):
Review Comment:
Optional: the other half of the issue is an operator that declares result as
a required positional argument, where the current code raises a TypeError
instead of silently passing None. Parametrizing this test over both signatures
— (self, context, result) and (self, context, result=None) — would cover that
crash path as well, at the cost of a single @pytest.mark.parametrize.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -6827,3 +6827,28 @@ def
test_stats_tags_with_standalone_and_key_value_tags(create_runtime_ti):
"task_id": "t",
"run_type": "manual",
}
+
+
+def test_post_execute_receives_result(create_runtime_ti,
mock_supervisor_comms):
+ """Verify that an overridden post_execute receives the result of the
execute method."""
+ post_execute_results: list[str] = []
+
+ from airflow.sdk import BaseOperator
Review Comment:
Both names are already imported at module scope — BaseOperator at line 58
and mock at line 31 — so these local imports can be removed. The project
convention is to keep imports at the top of the file. Also, the other run(...)
tests here use plain mock.MagicMock() for log; the spec=[...] is only needed by
the neighbouring test that exercises log.info.side_effect.
--
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]
