tests/test_chapter6.py

import sys
import unittest
from pathlib import Path

from langchain_core.messages import AIMessage, HumanMessage

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

from src.agents.agents import DietarySpecialist, OutputMerger


class StubLLM:
    def __init__(self, content):
        self.content = content
        self.calls = []

    async def ainvoke(self, messages, config=None):
        self.calls.append(messages)
        return AIMessage(content=self.content)

    async def astream(self, messages, config=None):
        self.calls.append(messages)
        yield AIMessage(content=self.content)


class TestChapter6(unittest.IsolatedAsyncioTestCase):
    async def test_output_merger_uses_latest_specialist_output(self):
        llm = StubLLM("merged response")
        merger = OutputMerger()
        merger.llm = llm

        state = {
            "messages": [
                HumanMessage(content="What should I do next?"),
                AIMessage(content="Older specialist response"),
                AIMessage(content="Older user-facing answer"),
            ],
            "clinician_outputs": [
                "Diagnosis says continue monitoring.",
                "Treatment says adjust medication dose.",
            ],
        }

        result = await merger.run(state)

        self.assertEqual(result["messages"][-1].content, "merged response")
        self.assertEqual(result["metrics"][0]["agent"], "OutputMerger")

        merged_prompt = "\n".join(
            message.content for message in llm.calls[-1]
        )
        self.assertIn("Treatment says adjust medication dose.", merged_prompt)
        self.assertNotIn("Older specialist response", merged_prompt)

    async def test_dietary_specialist_returns_metrics(self):
        llm = StubLLM("dietary guidance")
        specialist = DietarySpecialist()
        specialist.llm = llm

        state = {
            "messages": [HumanMessage(content="Suggest a low-carb meal plan for today.")]
        }

        result = await specialist.run(state)

        self.assertEqual(result["messages"][-1].content, "dietary guidance")
        self.assertEqual(result["metrics"][0]["agent"], "DietarySpecialist")


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