guan404ming commented on code in PR #72183:
URL: https://github.com/apache/airflow/pull/72183#discussion_r3965635686
##########
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:
Good catch, replaced with a parametrized dag_maker test asserting skipped
task IDs.
##########
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:
Thanks, documented the require_approval precondition on both flags in
docstring and rst.
--
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]
