kaxil commented on code in PR #74686:
URL: https://github.com/apache/airflow/pull/74686#discussion_r4268151823
##########
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:
This was the same slip: the extra `@patch` was the continuation test's
decorator. Removed in fea53285404, which also puts it back on that test.
--
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]
