"""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