tests/test_study_buddy_agent.py

import json
from unittest.mock import MagicMock

import pytest

from app.agents.study_buddy_agent import StudyBuddyAgent


class AsyncIterator:
    def __init__(self, items):
        self.items = list(items)

    def __aiter__(self):
        return self

    async def __anext__(self):
        if not self.items:
            raise StopAsyncIteration
        return self.items.pop(0)


@pytest.mark.asyncio
async def test_pair_buddy_default_ignores_node_label_and_uses_normal_chat():
    agent = StudyBuddyAgent.__new__(StudyBuddyAgent)
    client = MagicMock()
    client.stream_complete.return_value = AsyncIterator(["hello"])
    agent._client = client

    tokens = []
    async for token in agent.generate_initial_question(
        node_label="Old Topic Node",
        chunks=[{"source": "paper.pdf", "text": "The paper studies attention heads."}],
        familiarity="graduate",
        student_profile="",
    ):
        tokens.append(token)

    assert "".join(tokens) == "hello"
    messages = client.stream_complete.call_args[0][0]
    serialized = json.dumps(messages)
    assert "Old Topic Node" not in serialized
    assert "master the topic" not in serialized
    assert "Socratic" not in serialized
    assert "normal research companion" in serialized


@pytest.mark.asyncio
async def test_feynman_mode_keeps_questioning_pipeline():
    agent = StudyBuddyAgent.__new__(StudyBuddyAgent)
    client = MagicMock()
    client.complete.return_value = '{"has_gap": true, "reasoning": "missing mechanism"}'
    agent._client = client

    tokens = []
    async for token in agent.evaluate_and_ask_next(
        node_label="Ignored Label",
        chunks=[{"source": "paper.pdf", "text": "Attention uses query-key dot products."}],
        familiarity="graduate",
        history=[],
        student_answer="It is just a dictionary.",
        feynman_mode=True,
    ):
        tokens.append(token)

    assert "".join(tokens).startswith("__FLAG__")


@pytest.mark.asyncio
async def test_pair_buddy_accepts_optional_web_context_as_tool_context():
    agent = StudyBuddyAgent.__new__(StudyBuddyAgent)
    client = MagicMock()
    client.stream_complete.return_value = AsyncIterator(["ok"])
    agent._client = client

    tokens = []
    async for token in agent.evaluate_and_ask_next(
        node_label="Ignored Label",
        chunks=[],
        familiarity="graduate",
        history=[],
        student_answer="What is the latest DeepSeek model?",
        web_context="WEB RESEARCH:\n[DeepSeek update](https://example.com) says a new model shipped.",
        context_tools_summary="web_search",
    ):
        tokens.append(token)

    assert "".join(tokens) == "ok"
    messages = client.stream_complete.call_args[0][0]
    serialized = json.dumps(messages)
    assert "OPTIONAL CONTEXT TOOLS USED: web_search" in serialized
    assert "WEB RESEARCH:" in serialized
    assert "SOURCE MATERIAL" not in serialized