tests/test_graph.py

"""tests/test_graph.py — routing tests for the LangGraph agent.

The verify_node-related helpers (_normalize, _get_all_tool_outputs) referenced
by an earlier version of this file were removed when verify_node was retired
in favour of the post-synthesis reliability pass (see agent/post_synthesis.py).
"""
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage

from agent.graph import should_continue, nudge_node, AgentState


def _state(messages, tool_round_count, nudge_fired: bool = False) -> AgentState:
    return {
        "ticker": "AAPL",
        "messages": messages,
        "tool_round_count": tool_round_count,
        "nudge_fired": nudge_fired,
        "brief": None,
        "brief_markdown": None,
        "synthesis_error": None,
    }


def _ai_with_tools():
    return AIMessage(
        content="",
        tool_calls=[{"id": "c1", "name": "get_financial_metrics", "args": {"ticker": "AAPL"}}],
    )


def _ai_done():
    return AIMessage(content="I have all the information I need.", tool_calls=[])


def _tool_msg(name: str, content: str = "ok") -> ToolMessage:
    return ToolMessage(content=content, name=name, tool_call_id=f"id_{name}")


# ── Cap-based routing (existing tests, renamed) ──────────────────────────────


def test_routes_to_synthesis_when_no_tool_calls():
    # Floor satisfied: filing + transcript both touched.
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("search_filing"),
        _tool_msg("search_transcript"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=3)
    assert should_continue(state) == "synthesis"


def test_routes_to_tools_when_tool_calls_under_cap():
    state = _state([HumanMessage(content="brief"), _ai_with_tools()], tool_round_count=3)
    assert should_continue(state) == "tools"


def test_routes_to_synthesis_at_cap():
    state = _state([HumanMessage(content="brief"), _ai_with_tools()], tool_round_count=10)
    assert should_continue(state) == "synthesis"


def test_routes_to_synthesis_above_cap():
    state = _state([HumanMessage(content="brief"), _ai_with_tools()], tool_round_count=11)
    assert should_continue(state) == "synthesis"


# ── Nudge / minimum-evidence floor ───────────────────────────────────────────


def test_routes_to_nudge_when_no_filing_or_transcript_evidence():
    """Agent stops, but filing + transcript never called → force one more round."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("get_financial_metrics"),
        _tool_msg("get_analyst_expectations"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=2, nudge_fired=False)
    assert should_continue(state) == "nudge"


def test_routes_to_nudge_when_only_filing_present():
    """Filing touched but transcript missing → still nudge."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("search_filing"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=2, nudge_fired=False)
    assert should_continue(state) == "nudge"


def test_routes_to_synthesis_after_nudge_fired():
    """Even with missing coverage, only nudge once. Then trust the agent."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("get_financial_metrics"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=3, nudge_fired=True)
    assert should_continue(state) == "synthesis"


def test_routes_to_synthesis_when_filing_and_transcript_present():
    """Floor satisfied → no nudge needed, go to synthesis."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("search_filing"),
        _tool_msg("search_transcript"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=4, nudge_fired=False)
    assert should_continue(state) == "synthesis"


def test_nudge_not_triggered_when_at_cap():
    """At cap, even if floor unmet, route to synthesis (cap wins over floor)."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("get_financial_metrics"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=10, nudge_fired=False)
    assert should_continue(state) == "synthesis"


def test_nudge_node_sets_flag_and_appends_message():
    """nudge_node must set nudge_fired=True and inject a HumanMessage."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("get_financial_metrics"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=2, nudge_fired=False)
    result = nudge_node(state)

    assert result["nudge_fired"] is True
    assert len(result["messages"]) == 1
    new_msg = result["messages"][0]
    assert isinstance(new_msg, HumanMessage)
    # Should mention both missing tools when neither was called.
    assert "search_filing" in new_msg.content
    assert "search_transcript" in new_msg.content


def test_nudge_node_mentions_only_missing_tool():
    """If only one of {filing, transcript} is missing, nudge mentions just that one."""
    messages = [
        HumanMessage(content="brief"),
        _tool_msg("search_filing"),
        _ai_done(),
    ]
    state = _state(messages, tool_round_count=3, nudge_fired=False)
    result = nudge_node(state)

    new_msg = result["messages"][0]
    assert "search_transcript" in new_msg.content
    assert "search_filing" not in new_msg.content