vatsrahul1001 commented on code in PR #74686:
URL: https://github.com/apache/airflow/pull/74686#discussion_r4267985592
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_approval.py:
##########
@@ -398,6 +447,146 @@ def
test_a_second_approval_in_the_same_try_fails_closed(self):
upsert.assert_not_called()
assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data
+ @patch(REPEAT_SUPPORTED, True)
+ def test_a_later_step_pauses_again_with_its_own_call(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ first = _pause(_operator(_two_refunds, [shop.toolset()]),
_context(store))
+
+ with patch(UPSERT, autospec=True) as upsert:
+ with pytest.raises(TaskAwaitingInput) as second:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **first.kwargs, event=APPROVE
+ )
+
+ assert shop.calls == [("refund", 1)]
+ assert first.kwargs["pause_number"] == 1
+ assert second.value.kwargs["pause_number"] == 2
+ assert second.value.kwargs["tool_call_ids"] == ["c-2"]
+ assert upsert.call_args.kwargs["subject"] == "Approve refund for task
`t` (request 2)"
+ body = upsert.call_args.kwargs["body"]
+ assert '"order_id": 2' in body
+ assert '"order_id": 1' not in body
+ # The new pause's transcript, which its resume needs, survives the
resume that paused again.
+ assert _TOOL_APPROVAL_TRANSCRIPT_KEY in store.data
+
+ @patch(REPEAT_SUPPORTED, True)
+ def
test_an_agent_that_needs_ten_approvals_pauses_ten_times_and_finishes(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ model = _refunds_one_per_step(10)
+ paused = _pause(_operator(model, [shop.toolset()]), _context(store))
+ pauses = 1
+ output = None
+ while output is None:
+ ctx = _context(store)
+ with patch(UPSERT, autospec=True):
+ try:
+ output = _operator(model,
[shop.toolset()]).resume_after_tool_approval(
+ ctx, **paused.kwargs, event=APPROVE
+ )
+ except TaskAwaitingInput as again:
+ paused, pauses = again, pauses + 1
+
+ assert pauses == 10
+ # Every resume gets its own run_id; the last one is what the run_id
XCom holds.
+ ctx["task_instance"].xcom_push.assert_any_call(key="run_id",
value="ti-1-resumed-10")
+ assert shop.calls == [("refund", n) for n in range(1, 11)]
+ assert output == " | ".join(f"refunded order {n}" for n in range(1,
11))
+ assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data
+
+ @patch(REPEAT_SUPPORTED, True)
+ def test_a_rejected_later_step_carries_on_without_that_call(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ first = _pause(_operator(_two_refunds, [shop.toolset()]),
_context(store))
+ with patch(UPSERT, autospec=True):
+ with pytest.raises(TaskAwaitingInput) as second:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **first.kwargs, event=APPROVE
+ )
+
+ output = _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **second.value.kwargs, event=_reject("order 2 is
outside policy")
+ )
+
+ assert shop.calls == [("refund", 1)]
+ assert "refunded order 1" in output
+ assert "order 2 is outside policy" in output
+
+ @patch(REPEAT_SUPPORTED, True)
+ @pytest.mark.parametrize(
+ "stale_value",
+ [
+ pytest.param(1, id="number"),
+ pytest.param("1", id="string"),
+ pytest.param({"value": 1, "schema": {"type": "integer", "const":
1}}, id="param-spec"),
+ pytest.param("not-a-number", id="unreadable"),
+ ],
+ )
+ @patch(REPEAT_SUPPORTED, True)
+ def
test_an_answer_from_a_form_left_open_on_an_earlier_pause_asks_again(self,
stale_value):
+ """Reviewer B still has request 1's form open; their Approve must not
run refund 2."""
+ shop, store = _Shop(), _FakeTaskStateStore()
+ first = _pause(_operator(_two_refunds, [shop.toolset()]),
_context(store))
+ with patch(UPSERT, autospec=True):
+ with pytest.raises(TaskAwaitingInput) as second:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **first.kwargs, event=APPROVE
+ )
+ stale = {**APPROVE, "params_input": {"request": stale_value}}
+
+ with patch(UPSERT, autospec=True) as upsert:
+ with pytest.raises(TaskAwaitingInput) as asked_again:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **second.value.kwargs, event=stale
+ )
+
+ assert shop.calls == [("refund", 1)]
+ assert asked_again.value.kwargs["tool_call_ids"] == ["c-2"]
+ assert asked_again.value.kwargs["pause_number"] == 3
+ assert asked_again.value.kwargs["transcript_sha256"] ==
second.value.kwargs["transcript_sha256"]
+ assert '"order_id": 2' in upsert.call_args.kwargs["body"]
+ assert _TOOL_APPROVAL_TRANSCRIPT_KEY in store.data
+
+ output = _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **asked_again.value.kwargs, event={**APPROVE,
"params_input": {"request": 3}}
+ )
+ assert shop.calls == [("refund", 1), ("refund", 2)]
+ assert output == "refunded order 1 | refunded order 2"
+
+ @patch(REPEAT_SUPPORTED, True)
+ def test_an_answer_for_the_waiting_request_is_applied(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ first = _pause(_operator(_two_refunds, [shop.toolset()]),
_context(store))
+ with patch(UPSERT, autospec=True):
+ with pytest.raises(TaskAwaitingInput) as second:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **first.kwargs, event={**APPROVE,
"params_input": {"request": 1}}
+ )
+
+ output = _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **second.value.kwargs, event={**APPROVE,
"params_input": {"request": 2}}
+ )
+
+ assert output == "refunded order 1 | refunded order 2"
+
+ def
test_a_continuation_without_pause_number_resumes_as_the_first_pause(self):
Review Comment:
This one is missing `@patch(REPEAT_SUPPORTED, True)`, unlike the other
multi-pause tests in this class. It uses `_two_refunds`, so on the Airflow
3.3.2 compat job the real gate is False and the second pause hits
`ToolApprovalAlreadyRequestedError`. That's what's failing in `Compat
3.3.2:P3.11`. It passes on main and 3.4, which is why it only shows up there.
Adding the patch (or a skipif for < 3.4) should fix it.
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent_tool_approval.py:
##########
@@ -398,6 +447,146 @@ def
test_a_second_approval_in_the_same_try_fails_closed(self):
upsert.assert_not_called()
assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data
+ @patch(REPEAT_SUPPORTED, True)
+ def test_a_later_step_pauses_again_with_its_own_call(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ first = _pause(_operator(_two_refunds, [shop.toolset()]),
_context(store))
+
+ with patch(UPSERT, autospec=True) as upsert:
+ with pytest.raises(TaskAwaitingInput) as second:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **first.kwargs, event=APPROVE
+ )
+
+ assert shop.calls == [("refund", 1)]
+ assert first.kwargs["pause_number"] == 1
+ assert second.value.kwargs["pause_number"] == 2
+ assert second.value.kwargs["tool_call_ids"] == ["c-2"]
+ assert upsert.call_args.kwargs["subject"] == "Approve refund for task
`t` (request 2)"
+ body = upsert.call_args.kwargs["body"]
+ assert '"order_id": 2' in body
+ assert '"order_id": 1' not in body
+ # The new pause's transcript, which its resume needs, survives the
resume that paused again.
+ assert _TOOL_APPROVAL_TRANSCRIPT_KEY in store.data
+
+ @patch(REPEAT_SUPPORTED, True)
+ def
test_an_agent_that_needs_ten_approvals_pauses_ten_times_and_finishes(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ model = _refunds_one_per_step(10)
+ paused = _pause(_operator(model, [shop.toolset()]), _context(store))
+ pauses = 1
+ output = None
+ while output is None:
+ ctx = _context(store)
+ with patch(UPSERT, autospec=True):
+ try:
+ output = _operator(model,
[shop.toolset()]).resume_after_tool_approval(
+ ctx, **paused.kwargs, event=APPROVE
+ )
+ except TaskAwaitingInput as again:
+ paused, pauses = again, pauses + 1
+
+ assert pauses == 10
+ # Every resume gets its own run_id; the last one is what the run_id
XCom holds.
+ ctx["task_instance"].xcom_push.assert_any_call(key="run_id",
value="ti-1-resumed-10")
+ assert shop.calls == [("refund", n) for n in range(1, 11)]
+ assert output == " | ".join(f"refunded order {n}" for n in range(1,
11))
+ assert _TOOL_APPROVAL_TRANSCRIPT_KEY not in store.data
+
+ @patch(REPEAT_SUPPORTED, True)
+ def test_a_rejected_later_step_carries_on_without_that_call(self):
+ shop, store = _Shop(), _FakeTaskStateStore()
+ first = _pause(_operator(_two_refunds, [shop.toolset()]),
_context(store))
+ with patch(UPSERT, autospec=True):
+ with pytest.raises(TaskAwaitingInput) as second:
+ _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **first.kwargs, event=APPROVE
+ )
+
+ output = _operator(_two_refunds,
[shop.toolset()]).resume_after_tool_approval(
+ _context(store), **second.value.kwargs, event=_reject("order 2 is
outside policy")
+ )
+
+ assert shop.calls == [("refund", 1)]
+ assert "refunded order 1" in output
+ assert "order 2 is outside policy" in output
+
+ @patch(REPEAT_SUPPORTED, True)
+ @pytest.mark.parametrize(
+ "stale_value",
+ [
+ pytest.param(1, id="number"),
+ pytest.param("1", id="string"),
+ pytest.param({"value": 1, "schema": {"type": "integer", "const":
1}}, id="param-spec"),
+ pytest.param("not-a-number", id="unreadable"),
+ ],
+ )
+ @patch(REPEAT_SUPPORTED, True)
Review Comment:
Nit: `@patch(REPEAT_SUPPORTED, True)` is stacked twice here (above and below
the `parametrize`). It's harmless, because `patch` with an explicit value
injects no argument, but one of them can go.
--
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]
