| 1 | from __future__ import annotations
|
| 2 |
|
| 3 | import unittest
|
| 4 | from unittest.mock import patch
|
| 5 | from types import SimpleNamespace
|
| 6 |
|
| 7 | from tau2.data_model.message import AssistantMessage, SystemMessage, UserMessage
|
| 8 |
|
| 9 | from openclasp_agent import (
|
| 10 | CASE_RUNS,
|
| 11 | CaseRun,
|
| 12 | OpenClaspAgentState,
|
| 13 | OpenClaspReviewAgent,
|
| 14 | close_shield_cases,
|
| 15 | )
|
| 16 |
|
| 17 |
|
| 18 | class GenericReviewTests(unittest.TestCase):
|
| 19 | def test_anthropic_review_has_a_non_system_message(self):
|
| 20 | captured = []
|
| 21 |
|
| 22 | def fake_generate(**kwargs):
|
| 23 | captured.extend(kwargs["messages"])
|
| 24 | return AssistantMessage(role="assistant", content="Review result", cost=0.01)
|
| 25 |
|
| 26 | agent = OpenClaspReviewAgent(
|
| 27 | tools=[],
|
| 28 | domain_policy="Visible policy",
|
| 29 | llm="claude-sonnet-5",
|
| 30 | llm_args={},
|
| 31 | task_id="0",
|
| 32 | simulation_id="simulation-id",
|
| 33 | mode="generic-review",
|
| 34 | )
|
| 35 | state = OpenClaspAgentState(
|
| 36 | system_messages=[SystemMessage(role="system", content=agent.system_prompt)],
|
| 37 | messages=[UserMessage(role="user", content="Please cancel it")],
|
| 38 | simulation_id="simulation-id",
|
| 39 | )
|
| 40 |
|
| 41 | with patch("openclasp_agent.generate", side_effect=fake_generate):
|
| 42 | result, cost = agent._generic_review(state, '{"type":"user_response"}')
|
| 43 |
|
| 44 | self.assertEqual(result, "Review result")
|
| 45 | self.assertEqual(cost, 0.01)
|
| 46 | self.assertTrue(any(isinstance(message, UserMessage) for message in captured))
|
| 47 |
|
| 48 |
|
| 49 | class ShieldCaseCleanupTests(unittest.TestCase):
|
| 50 | def setUp(self):
|
| 51 | CASE_RUNS.clear()
|
| 52 |
|
| 53 | def tearDown(self):
|
| 54 | CASE_RUNS.clear()
|
| 55 |
|
| 56 | def test_closes_retry_cases_when_tau_returns_an_infra_error_placeholder(self):
|
| 57 | CASE_RUNS["real-attempt-id"] = CaseRun(
|
| 58 | case_id="case-id",
|
| 59 | task_id="39",
|
| 60 | consultation_count=1,
|
| 61 | )
|
| 62 | calls = []
|
| 63 |
|
| 64 | class FakeClient:
|
| 65 | def close_case(self, **kwargs):
|
| 66 | calls.append(kwargs)
|
| 67 | return {"outcomeId": "outcome-id"}
|
| 68 |
|
| 69 | results = SimpleNamespace(
|
| 70 | simulations=[
|
| 71 | SimpleNamespace(
|
| 72 | id="infra-placeholder-id",
|
| 73 | task_id="39",
|
| 74 | reward_info=None,
|
| 75 | termination_reason=SimpleNamespace(value="infrastructure_error"),
|
| 76 | )
|
| 77 | ]
|
| 78 | )
|
| 79 |
|
| 80 | closed = close_shield_cases(results, FakeClient())
|
| 81 |
|
| 82 | self.assertEqual(len(closed), 1)
|
| 83 | self.assertEqual(closed[0]["simulationId"], "real-attempt-id")
|
| 84 | self.assertEqual(closed[0]["result"], "unknown")
|
| 85 | self.assertEqual(calls[0]["result"], "unknown")
|
| 86 | self.assertEqual(CASE_RUNS, {})
|
| 87 |
|
| 88 |
|
| 89 | if __name__ == "__main__":
|
| 90 | unittest.main()
|