"""Offline regression tests. Run: python -m unittest -v test_agent_reliability""" import asyncio import io import json from pathlib import Path import tempfile import unittest from pydantic import ValidationError from agent_reliability_demo import ( BusinessError, Failure, Success, MultiplyOutput, Tool, close_logger, execute, make_logger, tool_schema, ) class ReliabilityTests(unittest.IsolatedAsyncioTestCase): def setUp(self): self.stream = io.StringIO() self.logger = make_logger(stream=self.stream) self.delays = [] def tearDown(self): close_logger(self.logger) async def sleep(self, delay): self.delays.append(delay) async def run_tool(self, payload=None, **options): return await execute( "multiply", {"a": 6, "b": 7} if payload is None else payload, logger=self.logger, run_id="run-demo", tool_call_id="call-demo", sleep=self.sleep, rng=lambda: 0.5, **options ) def events(self): return [json.loads(line) for line in self.stream.getvalue().splitlines()] async def test_success_and_identifiers(self): result = await self.run_tool() self.assertEqual(result.model_dump(), {"success": True, "data": {"value": 42}, "error": None}) event = self.events()[0] self.assertEqual((event["run_id"], event["tool_call_id"], event["attempt"]), ("run-demo", "call-demo", 1)) self.assertGreaterEqual(event["duration_ms"], 0) async def test_bad_inputs_never_invoke_tool(self): calls = [] async def tool(args): calls.append(args) return {"value": 42} for payload in [{"a": "6", "b": 7}, {"a": True, "b": 7}, {"a": 1}, {"a": 1, "b": 2, "token": "secret-example"}, {"a": 10001, "b": 2}]: with self.subTest(payload=payload): result = await self.run_tool(payload, registry={"multiply": Tool(tool)}) self.assertEqual(result.error, "input_validation") self.assertEqual(calls, []) self.assertNotIn("secret-example", self.stream.getvalue()) async def test_unknown_tool(self): result = await self.run_tool(registry={}) self.assertEqual(result.error, "unknown_tool") self.assertEqual(self.events()[0]["attempt"], 0) async def test_output_failure_no_retry(self): calls = [] async def broken(args): calls.append(1) return {"value": "42"} result = await self.run_tool(registry={"multiply": Tool(broken, True)}) self.assertEqual(result.error, "output_validation") self.assertEqual(len(calls), 1) async def test_business_and_internal_no_retry(self): for exception, expected in [(BusinessError, "business"), (ValueError, "internal")]: calls = [] async def broken(args): calls.append(1) raise exception("private exception body") result = await self.run_tool(registry={"multiply": Tool(broken, True)}) self.assertEqual(result.error, expected) self.assertEqual(len(calls), 1) self.assertNotIn("private exception body", self.stream.getvalue()) async def test_connection_recovers_with_jitter(self): calls = [] async def flaky(args): calls.append(1) if len(calls) < 3: raise ConnectionError() return {"value": 42} result = await self.run_tool(registry={"multiply": Tool(flaky, True)}) self.assertTrue(result.success) self.assertEqual(self.delays, [0.025, 0.05]) self.assertEqual([e["attempt"] for e in self.events()], [1, 2, 3]) async def test_retry_exhaustion_and_cap(self): async def broken(args): raise ConnectionError() result = await self.run_tool(registry={"multiply": Tool(broken, True)}, base_delay=1, cap_delay=0.1) self.assertEqual(result.error, "connection") self.assertEqual(self.delays, [0.05, 0.05]) self.assertEqual(len(self.events()), 3) async def test_unsafe_tool_not_retried(self): async def broken(args): raise ConnectionError() result = await self.run_tool(registry={"multiply": Tool(broken)}) self.assertFalse(result.success) self.assertEqual(len(self.events()), 1) self.assertEqual(self.delays, []) async def test_cooperative_timeout_and_cleanup(self): cleanup = [] async def slow(args): try: await asyncio.sleep(30) finally: cleanup.append(True) result = await self.run_tool(registry={"multiply": Tool(slow)}, timeout_s=0.01) self.assertEqual(result.error, "timeout") self.assertEqual(cleanup, [True]) async def test_external_cancellation_propagates(self): started = asyncio.Event() async def pending(args): started.set() await asyncio.sleep(30) task = asyncio.create_task(self.run_tool(registry={"multiply": Tool(pending, True)})) await started.wait() task.cancel() with self.assertRaises(asyncio.CancelledError): await task self.assertEqual(self.events()[-1]["error_type"], "cancelled") self.assertEqual(self.delays, []) async def test_invalid_runtime_options(self): for options in [{"max_attempts": 0}, {"max_attempts": True}, {"max_attempts": 6}, {"timeout_s": 0}, {"timeout_s": float("nan")}, {"base_delay": -1}]: with self.subTest(options=options), self.assertRaises(ValueError): await self.run_tool(**options) async def test_rotation_keeps_bounded_files(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "exercise.log" logger = make_logger(path, io.StringIO(), max_bytes=250, backup_count=2) try: for i in range(10): await execute("multiply", {"a": 1, "b": 2}, logger=logger, run_id="r", tool_call_id=str(i)) finally: close_logger(logger) self.assertEqual(len(list(Path(directory).glob("exercise.log*"))), 3) for file in Path(directory).glob("exercise.log*"): for line in file.read_text(encoding="utf-8").splitlines(): self.assertEqual(json.loads(line)["event"], "tool_attempt") async def test_schema_matches_contract(self): schema = tool_schema()["parameters"] self.assertEqual(schema["required"], ["a", "b"]) self.assertFalse(schema["additionalProperties"]) async def test_result_envelope_rejects_contradictions(self): with self.assertRaises(ValidationError): Success(data=MultiplyOutput(value=42), error="timeout") with self.assertRaises(ValidationError): Failure(error="timeout", data={"value": 42}) if __name__ == "__main__": unittest.main()