tests/test_agent_output_structuring.py

import asyncio
import sys
import types
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.agents.agents import ResponseValidator, SafetyCheck


class _DummyLLM:
    def __init__(self, content):
        self.content = content

    async def ainvoke(self, messages):
        return types.SimpleNamespace(
            content=self.content,
            usage_metadata=None,
            response_metadata={}
        )


class TestStructuredAgentOutputs(unittest.TestCase):
    def test_response_validator_accepts_structured_json(self):
        validator = ResponseValidator()
        validator.llm = _DummyLLM('{"decision": "invalid", "reason": "not complete"}')

        async def _run():
            return await validator.run({"messages": [types.SimpleNamespace(content="Sample output")]})

        result = asyncio.run(_run())
        self.assertFalse(result["is_valid"])

    def test_safety_check_accepts_structured_json(self):
        safety_check = SafetyCheck()
        safety_check.llm = _DummyLLM('{"decision": "safe"}')

        async def _run():
            return await safety_check.run({"messages": [types.SimpleNamespace(content="Sample output")]})

        result = asyncio.run(_run())
        self.assertTrue(result["is_safe"])


if __name__ == "__main__":
    unittest.main()