kaxil commented on code in PR #72183:
URL: https://github.com/apache/airflow/pull/72183#discussion_r3963568514
##########
providers/common/ai/tests/unit/common/ai/operators/test_llm_branch.py:
##########
@@ -420,6 +420,29 @@ def
test_execute_complete_reject_skips_downstream_except_teardowns(self, mock_do
assert list(mock_skip.call_args.kwargs["tasks"]) == [task_a]
mock_do_branch.assert_not_called()
+ @patch.object(LLMBranchOperator, "skip")
+ @patch.object(LLMBranchOperator, "do_branch")
+ def
test_execute_complete_reject_skips_indirect_downstream_when_ignoring_trigger_rules(
+ self, mock_do_branch, mock_skip
+ ):
+ op = LLMBranchOperator(task_id="t", prompt="p", llm_conn_id="c",
ignore_downstream_trigger_rules=True)
+ op.downstream_task_ids = {"task_a"}
+ event = {"chosen_options": ["Reject"], "responded_by_user": "admin"}
+ task_a = MagicMock(is_teardown=False)
+ indirect = MagicMock(is_teardown=False)
+ cleanup = MagicMock(is_teardown=True)
+ task = MagicMock()
+ task.get_flat_relatives.return_value = [task_a, indirect, cleanup]
+ ti = MagicMock()
+ ctx = MagicMock(**{"__getitem__": lambda self, key: {"task": task,
"ti": ti}[key]})
+
+ op.execute_complete(ctx, generated_output="task_a", event=event)
+
+ task.get_flat_relatives.assert_called_once_with(upstream=False)
+ task.get_direct_relatives.assert_not_called()
+ assert list(mock_skip.call_args.kwargs["tasks"]) == [task_a, indirect]
Review Comment:
These assert which getter was called and echo back the stubbed return value,
so the traversal itself is never exercised -- `get_flat_relatives` is a
`MagicMock` here, and the test passes whatever a real Dag would actually return.
`TestShortCircuitWithTeardown::test_short_circuit_with_teardowns` in
`providers/standard/tests/unit/standard/operators/test_python.py` covers the
same flag against a real Dag: `dag_maker` with `op1 >> op2 >> op3 >> op4`,
parametrized over `ignore_downstream_trigger_rules` and `with_teardown`,
asserting the skipped task IDs. The split there between a teardown and no
teardown -- `["op2", "op3"]` against `["op2", "op3", "op4"]` -- is what pins
the transitive-plus-teardown behaviour that a mocked return value can't reach.
Since the teardown carve-out in #71073 only surfaced on a real Dag, one
`dag_maker` case here would carry more than the two call assertions.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm_branch.py:
##########
@@ -78,12 +83,14 @@ def __init__(
*,
allow_multiple_branches: bool = False,
fail_on_reject: bool = False,
+ ignore_downstream_trigger_rules: bool = False,
Review Comment:
`execute_complete` is only reachable through `defer_for_approval`, which
`execute` calls under `if self.require_approval`. So this flag does nothing at
all unless `require_approval=True`, and neither the param docs nor the rst
mention that precondition.
That is also where the `ApprovalOperator` parity the docs claim stops
holding: there the operator always defers, so the flag always applies. Someone
carrying the idiom across would set it, see no change in behaviour, and have
nothing pointing at why.
`self.require_approval` is already set by the `super().__init__()` just
below, so raising here is possible -- though it would mean the existing reject
tests need `require_approval=True` to construct, so a sentence on the param is
probably the cheaper fix. Same property applies to `fail_on_reject`, so this
could ride along with moving that one to `LLMApprovalMixin`.
--
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]
