kaxil commented on code in PR #74686:
URL: https://github.com/apache/airflow/pull/74686#discussion_r4268139270
##########
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:
Good catch, thanks. When I added the stale-answer tests above it, its
`@patch(REPEAT_SUPPORTED, True)` ended up on the new parametrized test instead.
Fixed in fea53285404. I also ran the whole file with the gate forced closed, as
on 3.3, to check that no other multi-pause test relies on the real gate.
--
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]
