tests/agents/test_intake.py

"""Tests for agents/intake.py.

_months_from_date and _resolve_months: pure date helpers (freezegun).
stream_intake_turn: FakeStream pattern for the streaming generator.
"""
from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest
from freezegun import freeze_time

from agents.intake import _months_from_date, _resolve_months, stream_intake_turn
from models import PatientProfile
from tests.conftest import FakeStream, make_message, make_text_block, make_tool_use_block


# ---------------------------------------------------------------------------
# _months_from_date
# ---------------------------------------------------------------------------

class TestMonthsFromDate:
    @freeze_time("2026-05-01")
    def test_same_month_is_zero(self):
        assert _months_from_date("2026-05") == 0

    @freeze_time("2026-05-01")
    def test_one_month_ago(self):
        assert _months_from_date("2026-04") == 1

    @freeze_time("2026-05-01")
    def test_twelve_months_ago(self):
        assert _months_from_date("2025-05") == 12

    @freeze_time("2026-05-01")
    def test_crosses_year_boundary(self):
        assert _months_from_date("2025-11") == 6

    @freeze_time("2026-05-01")
    def test_multi_year(self):
        assert _months_from_date("2024-05") == 24

    def test_invalid_string_returns_zero(self):
        assert _months_from_date("not-a-date") == 0

    def test_empty_string_returns_zero(self):
        assert _months_from_date("") == 0

    @freeze_time("2026-05-01")
    def test_future_date_returns_negative(self):
        # The function does not clamp — a future onset_date from the LLM produces
        # a negative onset_months, which can cause eligibility criteria like
        # "onset_months >= 6" to pass incorrectly. Documented here so a clamping
        # fix is validated when added.
        result = _months_from_date("2030-01")
        assert result < 0


# ---------------------------------------------------------------------------
# _resolve_months
# ---------------------------------------------------------------------------

class TestResolveMonths:
    @freeze_time("2026-05-01")
    def test_prefers_date_key(self):
        data = {"onset_date": "2025-05", "onset_months": 99}
        assert _resolve_months(data, "onset_date", "onset_months") == 12

    def test_falls_back_to_months_key(self):
        data = {"onset_months": 18}
        assert _resolve_months(data, "onset_date", "onset_months") == 18

    def test_missing_both_returns_zero(self):
        assert _resolve_months({}, "onset_date", "onset_months") == 0

    def test_none_months_returns_zero(self):
        data = {"onset_months": None}
        assert _resolve_months(data, "onset_date", "onset_months") == 0


# ---------------------------------------------------------------------------
# stream_intake_turn — FakeStream pattern
# ---------------------------------------------------------------------------

def _initial_messages() -> list:
    return [
        {"role": "user", "content": "Please begin."},
        {"role": "assistant", "content": [make_text_block("Hello, I'm Beacon.")]},
        {"role": "user", "content": "I have ALS, age 52."},
    ]


class TestStreamIntakeTurnText:
    def test_yields_token_events(self, mock_client):
        final_msg = make_message(
            content=[make_text_block("What is your ZIP code?")],
            stop_reason="end_turn",
        )
        mock_client.messages.stream.return_value = FakeStream(
            tokens=["What ", "is ", "your ", "ZIP?"],
            final_message=final_msg,
        )

        events = list(stream_intake_turn(mock_client, _initial_messages()))
        token_events = [e for e in events if e[0] == "token"]
        assert token_events == [
            ("token", "What "), ("token", "is "), ("token", "your "), ("token", "ZIP?")
        ]

    def test_yields_text_event_with_updated_messages(self, mock_client):
        final_msg = make_message(
            content=[make_text_block("What is your ZIP code?")],
            stop_reason="end_turn",
        )
        mock_client.messages.stream.return_value = FakeStream(
            tokens=["What is your ZIP code?"],
            final_message=final_msg,
        )

        events = list(stream_intake_turn(mock_client, _initial_messages()))
        text_events = [e for e in events if e[0] == "text"]
        assert len(text_events) == 1
        kind, text, msgs = text_events[0]
        assert "ZIP" in text
        assert isinstance(msgs, list)

    def test_single_stream_call_for_text_response(self, mock_client):
        final_msg = make_message(
            content=[make_text_block("OK")], stop_reason="end_turn"
        )
        mock_client.messages.stream.return_value = FakeStream(tokens=["OK"], final_message=final_msg)
        list(stream_intake_turn(mock_client, _initial_messages()))
        assert mock_client.messages.stream.call_count == 1


class TestStreamIntakeTurnIdentifyDisease:
    def test_identify_disease_yields_reset_stream_and_continues(self, mock_client):
        identify_block = make_tool_use_block(
            "identify_disease",
            {"standardized_name": "Amyotrophic Lateral Sclerosis"},
            tool_use_id="tu_id1",
        )
        # First stream: identify_disease tool_use
        first_msg = make_message(content=[identify_block], stop_reason="tool_use")
        # Second stream: text response after tool result
        second_msg = make_message(
            content=[make_text_block("Got it, collecting ALS benchmarks.")],
            stop_reason="end_turn",
        )
        mock_client.messages.stream.side_effect = [
            FakeStream(tokens=[], final_message=first_msg),
            FakeStream(tokens=["Got it"], final_message=second_msg),
        ]

        events = list(stream_intake_turn(mock_client, _initial_messages()))
        assert ("reset_stream",) in events
        assert mock_client.messages.stream.call_count == 2

    def test_messages_contain_tool_result_after_identify(self, mock_client):
        identify_block = make_tool_use_block(
            "identify_disease",
            {"standardized_name": "Amyotrophic Lateral Sclerosis"},
            tool_use_id="tu_id1",
        )
        first_msg = make_message(content=[identify_block], stop_reason="tool_use")
        second_msg = make_message(
            content=[make_text_block("Collecting ALSFRS-R.")], stop_reason="end_turn"
        )
        mock_client.messages.stream.side_effect = [
            FakeStream(tokens=[], final_message=first_msg),
            FakeStream(tokens=["Collecting"], final_message=second_msg),
        ]

        events = list(stream_intake_turn(mock_client, _initial_messages()))
        text_event = next(e for e in events if e[0] == "text")
        msgs = text_event[2]
        roles = [m["role"] for m in msgs]
        assert "user" in roles  # tool_result is in a user message


class TestStreamIntakeTurnSubmitProfile:
    def _submit_block(self) -> MagicMock:
        return make_tool_use_block(
            "submit_profile",
            {
                "disease": "Amyotrophic Lateral Sclerosis",
                "age": 52,
                "onset_date": "2024-11",
                "diagnosis_date": "2025-05",
                "zip_code": "02115",
                "country_code": "US",
                "radius_miles": 100,
                "phases": ["2", "3"],
                "include_eap": False,
                "include_observational": False,
            },
            tool_use_id="tu_sub1",
        )

    def test_yields_profile_event(self, mock_client):
        submit_block = self._submit_block()
        final_msg = make_message(content=[submit_block], stop_reason="tool_use")
        mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)

        with patch("agents.intake.geocode_zip", return_value=(42.33, -71.10)):
            events = list(stream_intake_turn(mock_client, _initial_messages()))

        profile_events = [e for e in events if e[0] == "profile"]
        assert len(profile_events) == 1

    def test_profile_event_contains_patient_profile(self, mock_client):
        submit_block = self._submit_block()
        final_msg = make_message(content=[submit_block], stop_reason="tool_use")
        mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)

        with patch("agents.intake.geocode_zip", return_value=(42.33, -71.10)):
            events = list(stream_intake_turn(mock_client, _initial_messages()))

        _, profile, msgs = next(e for e in events if e[0] == "profile")
        assert isinstance(profile, PatientProfile)
        assert profile.disease == "Amyotrophic Lateral Sclerosis"
        assert profile.age == 52

    def test_geocode_failure_sets_zero_coords(self, mock_client):
        submit_block = self._submit_block()
        final_msg = make_message(content=[submit_block], stop_reason="tool_use")
        mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)

        with patch("agents.intake.geocode_zip", side_effect=Exception("timeout")):
            events = list(stream_intake_turn(mock_client, _initial_messages()))

        _, profile, _ = next(e for e in events if e[0] == "profile")
        assert profile.lat == 0.0
        assert profile.lon == 0.0

    def test_generator_terminates_after_profile(self, mock_client):
        submit_block = self._submit_block()
        final_msg = make_message(content=[submit_block], stop_reason="tool_use")
        mock_client.messages.stream.return_value = FakeStream(tokens=[], final_message=final_msg)

        with patch("agents.intake.geocode_zip", return_value=(42.33, -71.10)):
            events = list(stream_intake_turn(mock_client, _initial_messages()))

        # After profile, no more events
        profile_idx = next(i for i, e in enumerate(events) if e[0] == "profile")
        assert profile_idx == len(events) - 1