tests/test_parser.py

import json
import unittest
from unittest.mock import patch

from training_coach.parser import (
    PARSER_FALLBACK_MODEL,
    PARSER_LOG_PARSED_JSON_ENV_VAR,
    PARSER_LOG_RESPONSES_ENV_VAR,
    PARSER_MODEL,
    apply_follow_up_triggers,
    build_parser_messages,
    expected_response_format,
    normalize_parsed_check_in,
    parse_model_response,
)
from training_coach.models import CheckIn, FollowUpQuestion, PainIssue, ParsedCheckIn


class ParserTest(unittest.TestCase):
    def test_parser_model_choice_is_locked(self):
        self.assertEqual(PARSER_MODEL, "Qwen/Qwen2.5-1.5B-Instruct")
        self.assertEqual(PARSER_FALLBACK_MODEL, "Qwen/Qwen3-4B")

    def test_build_parser_messages_includes_response_shape_and_raw_text(self):
        messages = build_parser_messages("45 min, terrible night, low energy")

        self.assertEqual(messages[0]["role"], "system")
        self.assertIn("Return JSON only", messages[0]["content"])
        self.assertIn("Do not rename fields or use synonyms", messages[0]["content"])
        self.assertIn("Do not ask post-workout questions", messages[0]["content"])
        self.assertIn("sleep_hours", messages[0]["content"])
        self.assertIn("follow_up_items", messages[0]["content"])
        self.assertIn("recent_endurance_run", messages[0]["content"])
        self.assertIn("triceps_brachii", messages[0]["content"])
        self.assertIn("exactly this shape", messages[1]["content"])
        self.assertIn("sleep_quality", messages[1]["content"])
        self.assertIn("pain_issues", messages[1]["content"])
        self.assertIn("context_signals", messages[1]["content"])
        self.assertIn('"gastrocnemius"', messages[1]["content"])
        self.assertIn("Do not add extra keys", messages[1]["content"])
        self.assertIn("45 min, terrible night, low energy", messages[1]["content"])
        self.assertTrue(messages[1]["content"].endswith("/no_think"))

    def test_expected_response_format_is_parsed_check_in_schema(self):
        schema = expected_response_format()

        self.assertEqual(schema["title"], "ParsedCheckIn")
        self.assertIn("check_in", schema["properties"])
        self.assertIn("follow_up_items", schema["properties"])
        self.assertIn("follow_up_questions", schema["properties"])
        self.assertIn("context_signals", schema["properties"])

    def test_parse_model_response_accepts_poor_sleep_follow_up(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "45 min, terrible night, low energy",
                    "time_available_minutes": 45,
                    "energy_level": "low",
                    "sleep_quality": "poor",
                    "sleep_hours": None,
                    "pain_or_injury": "unsure",
                },
                "missing_fields": ["sleep_hours"],
                "follow_up_items": [
                    {
                        "field": "sleep_hours",
                        "question": "How many hours did you sleep?",
                        "reason": "Sleep quality is poor but sleep duration is missing.",
                    }
                ],
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(parsed.check_in.sleep_quality, "poor")
        self.assertIsNone(parsed.check_in.sleep_hours)
        self.assertIn("sleep_hours", parsed.missing_fields)
        self.assertEqual(parsed.follow_up_questions, ["How many hours did you sleep?"])

    def test_parse_model_response_lifts_top_level_keys_misplaced_in_check_in(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "45 min, terrible night",
                    "time_available_minutes": 45,
                    "sleep_quality": "poor",
                    "pain_or_injury": "no",
                    "missing_fields": ["sleep_hours"],
                    "follow_up_items": [
                        {
                            "field": "sleep_hours",
                            "question": "How many hours did you sleep?",
                            "reason": "Sleep quality is poor but hours are missing.",
                        }
                    ],
                    "follow_up_questions": ["How many hours did you sleep?"],
                    "context_signals": [],
                    "notes": "check-in note stays here",
                }
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(parsed.check_in.time_available_minutes, 45)
        self.assertEqual(parsed.check_in.notes, "check-in note stays here")
        self.assertIn("sleep_hours", parsed.missing_fields)
        self.assertEqual(parsed.follow_up_questions, ["How many hours did you sleep?"])

    def test_parse_model_response_drops_unknown_keys_everywhere(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "45 min, tricep ache, ran a 10k yesterday",
                    "time_available_minutes": 45,
                    "pain_or_injury": "yes",
                    "pain_issues": [
                        {
                            "affected_muscle": "triceps_brachii",
                            "severity": "mild",
                            "notes": "tricep ache",
                            "side": "left",
                        }
                    ],
                    "warmup_done": True,
                },
                "training_recommendation": "go lighter today",
                "follow_up_items": [
                    {
                        "field": "sleep_hours",
                        "question": "How many hours did you sleep?",
                        "reason": "Sleep info missing.",
                        "priority": 1,
                    }
                ],
                "context_signals": [
                    {
                        "label": "recent_endurance_run",
                        "evidence": "ran a 10k",
                        "follow_up_question": "",
                        "confidence": 0.9,
                    }
                ],
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(parsed.check_in.time_available_minutes, 45)
        self.assertEqual(
            parsed.check_in.pain_issues[0].affected_muscle, "triceps_brachii"
        )
        self.assertEqual(parsed.follow_up_items[0].field, "sleep_hours")
        self.assertEqual(parsed.context_signals[0].label, "recent_endurance_run")

    def test_parse_model_response_coerces_invalid_nested_values(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "30 min, shoulder pain",
                    "time_available_minutes": -5,
                    "sleep_hours": 99,
                    "pain_or_injury": "yes",
                    "pain_issues": [
                        {
                            "affected_muscle": "left_shoulder",
                            "severity": "terrible",
                            "notes": "",
                        }
                    ],
                },
                "follow_up_items": [
                    {"field": "shoulder", "question": "Which part of the shoulder?", "reason": ""},
                    {"field": "energy_level", "question": "", "reason": "no question"},
                ],
                "context_signals": [
                    {"label": "ran_marathon", "evidence": "mentioned a run", "follow_up_question": ""},
                    {"label": "travel", "evidence": "", "follow_up_question": ""},
                ],
            }
        )

        parsed = parse_model_response(response)

        self.assertIsNone(parsed.check_in.time_available_minutes)
        self.assertIsNone(parsed.check_in.sleep_hours)
        issue = parsed.check_in.pain_issues[0]
        self.assertIsNone(issue.affected_muscle)
        self.assertEqual(issue.severity, "unsure")
        self.assertTrue(issue.notes)
        self.assertEqual(len(parsed.follow_up_items), 1)
        self.assertEqual(parsed.follow_up_items[0].field, "other")
        self.assertEqual(len(parsed.context_signals), 1)
        self.assertEqual(parsed.context_signals[0].label, "other")

    def test_parse_model_response_nulls_invalid_check_in_enums(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "30 min, stressed from work",
                    "time_available_minutes": 30,
                    "energy_level": "neutral",
                    "sleep_quality": "great",
                    "mood_stress": "stressed",
                    "pain_or_injury": "maybe",
                }
            }
        )

        parsed = parse_model_response(response)

        self.assertIsNone(parsed.check_in.energy_level)
        self.assertIsNone(parsed.check_in.sleep_quality)
        self.assertEqual(parsed.check_in.mood_stress, "stressed")
        self.assertEqual(parsed.check_in.pain_or_injury, "unsure")

    def test_parse_model_response_drops_bare_field_names_in_follow_up_questions(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "60 min, slept great, no pain",
                    "time_available_minutes": 60,
                    "sleep_quality": "good",
                    "pain_or_injury": "no",
                },
                "follow_up_questions": ["energy_level", "mood_stress"],
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(parsed.follow_up_questions, [])

    def test_parse_model_response_accepts_adjacent_context_signal(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "Yesterday I ran a 10k.",
                    "pain_or_injury": "unsure",
                },
                "context_signals": [
                    {
                        "label": "recent_endurance_run",
                        "evidence": "Yesterday I ran a 10k.",
                        "follow_up_question": (
                            "How are your legs and energy after yesterday's run?"
                        ),
                    }
                ],
                "follow_up_questions": [
                    "How are your legs and energy after yesterday's run?"
                ],
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(parsed.context_signals[0].label, "recent_endurance_run")
        self.assertEqual(
            parsed.context_signals[0].follow_up_question,
            "How are your legs and energy after yesterday's run?",
        )

    def test_parse_model_response_accepts_pain_issues_per_problem_area(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "Right tricep hurts and my hamstring is tight.",
                    "pain_or_injury": "yes",
                    "pain_issues": [
                        {
                            "affected_muscle": "triceps_brachii",
                            "severity": "moderate",
                            "notes": "Right tricep hurts.",
                        },
                        {
                            "affected_muscle": "hamstrings",
                            "severity": "mild",
                            "notes": "Hamstring is tight.",
                        },
                    ],
                },
                "follow_up_questions": [
                    "Is the hamstring tightness normal soreness or pain?"
                ],
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(len(parsed.check_in.pain_issues), 2)
        self.assertEqual(
            parsed.check_in.pain_issues[0].affected_muscle,
            "triceps_brachii",
        )

    def test_parse_model_response_fills_obvious_missing_affected_muscle(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": (
                        "My right tricep hurts today, probably moderate. "
                        "I have about an hour."
                    ),
                    "time_available_minutes": 60,
                    "pain_or_injury": "yes",
                    "pain_issues": [
                        {
                            "severity": "moderate",
                            "notes": "Right tricep hurts today.",
                        }
                    ],
                },
            }
        )

        parsed = parse_model_response(response)

        self.assertEqual(
            parsed.check_in.pain_issues[0].affected_muscle,
            "triceps_brachii",
        )

    def test_normalize_parsed_check_in_creates_pain_issue_from_bicep_tore(self):
        parsed = normalize_parsed_check_in(
            ParsedCheckIn(check_in=CheckIn(raw_text="My bicep tore."))
        )

        self.assertEqual(parsed.check_in.pain_or_injury, "yes")
        self.assertEqual(len(parsed.check_in.pain_issues), 1)
        self.assertEqual(
            parsed.check_in.pain_issues[0].affected_muscle,
            "biceps_brachii",
        )

    def test_normalize_removes_unsupported_run_follow_up_and_maps_glutes(self):
        parsed = normalize_parsed_check_in(
            ParsedCheckIn(
                check_in=CheckIn(
                    raw_text=(
                        "i feel tired, and i have 45 minutes to exercise. "
                        "My glutes are in pain, not recovered."
                    ),
                    time_available_minutes=45,
                    pain_or_injury="yes",
                    pain_issues=[
                        PainIssue(
                            affected_muscle=None,
                            severity="unsure",
                            notes="glutes are in pain, not recovered",
                        )
                    ],
                ),
                context_signals=[
                    {
                        "label": "recent_endurance_run",
                        "evidence": "not recovered",
                        "follow_up_question": "How are your legs and energy after the run?",
                    }
                ],
                follow_up_questions=[
                    "How are your legs and energy after the run?",
                    "You mentioned a body area. Is it pain, injury, normal soreness, or just a note?",
                ],
                follow_up_items=[
                    FollowUpQuestion(
                        field="context_signals",
                        question="How are your legs and energy after the run?",
                        reason="The model thought the user had run.",
                    ),
                    FollowUpQuestion(
                        field="pain_issues",
                        question=(
                            "You mentioned a body area. Is it pain, injury, "
                            "normal soreness, or just a note?"
                        ),
                        reason="The model thought the body area was unclear.",
                    ),
                ],
            )
        )

        self.assertEqual(parsed.context_signals, [])
        self.assertEqual(parsed.follow_up_items, [])
        self.assertEqual(parsed.follow_up_questions, [])
        self.assertEqual(
            parsed.check_in.pain_issues[0].affected_muscle,
            "gluteus_maximus",
        )

    def test_normalize_maps_each_pain_issue_from_its_own_notes(self):
        parsed = normalize_parsed_check_in(
            ParsedCheckIn(
                check_in=CheckIn(
                    raw_text=(
                        "my calves are cooked and in pain from last round. "
                        "my biceps are a bit in pain in contraction."
                    ),
                    pain_or_injury="yes",
                    pain_issues=[
                        PainIssue(
                            affected_muscle=None,
                            severity="unsure",
                            notes="calves are cooked and in pain from last round",
                        ),
                        PainIssue(
                            affected_muscle=None,
                            severity="unsure",
                            notes="biceps are a bit in pain in contraction",
                        ),
                    ],
                )
            )
        )

        self.assertEqual(
            parsed.check_in.pain_issues[0].affected_muscle,
            "gastrocnemius",
        )
        self.assertEqual(
            parsed.check_in.pain_issues[1].affected_muscle,
            "biceps_brachii",
        )

    def test_normalize_parsed_check_in_leaves_unclear_muscle_unknown(self):
        parsed = normalize_parsed_check_in(
            ParsedCheckIn(
                check_in=CheckIn(
                    raw_text="My arm feels weird.",
                    pain_issues=[
                        PainIssue(
                            affected_muscle=None,
                            severity="unsure",
                            notes="Arm feels weird.",
                        )
                    ],
                )
            )
        )

        self.assertIsNone(parsed.check_in.pain_issues[0].affected_muscle)

    def test_follow_up_triggers_keep_model_sleep_question_when_missing(self):
        parsed = apply_follow_up_triggers(
            ParsedCheckIn(
                check_in=CheckIn(raw_text="Rough night, maybe a nap later."),
                missing_fields=["sleep_hours"],
                follow_up_items=[
                    FollowUpQuestion(
                        field="sleep_hours",
                        question="How many hours did you sleep?",
                        reason="Sleep was rough and duration is missing.",
                    )
                ],
            )
        )

        self.assertIn("sleep_hours", parsed.missing_fields)
        self.assertEqual(parsed.follow_up_questions, ["How many hours did you sleep?"])

    def test_follow_up_triggers_keep_only_missing_sleep_hours_question(self):
        parsed = apply_follow_up_triggers(
            ParsedCheckIn(
                check_in=CheckIn(
                    raw_text="Sleep was meh.",
                    sleep_quality="okay",
                    sleep_hours=None,
                ),
                follow_up_items=[
                    FollowUpQuestion(
                        field="sleep_quality",
                        question="Was your sleep poor, okay, or good?",
                        reason="The model duplicated an answered sleep field.",
                    ),
                    FollowUpQuestion(
                        field="sleep_hours",
                        question="How many hours did you sleep?",
                        reason="Duration is still missing.",
                    ),
                ],
            )
        )

        self.assertEqual(parsed.follow_up_questions, ["How many hours did you sleep?"])

    def test_follow_up_triggers_do_not_ask_sleep_when_sleep_data_is_complete(self):
        parsed = apply_follow_up_triggers(
            ParsedCheckIn(
                check_in=CheckIn(
                    raw_text=(
                        "im ok, sleep was meh, and i have 30 minutos to exercise. "
                        "I slept 5 hours, it was okay."
                    ),
                    sleep_quality="okay",
                    sleep_hours=5,
                ),
                missing_fields=["sleep_hours"],
                follow_up_questions=[
                    "How many hours did you sleep?",
                    "How many hours did you sleep, and was the sleep poor, okay, or good?",
                ],
                follow_up_items=[
                    FollowUpQuestion(
                        field="sleep_hours",
                        question="How many hours did you sleep?",
                        reason="The model duplicated an answered sleep field.",
                    )
                ],
            )
        )

        self.assertNotIn("sleep_hours", parsed.missing_fields)
        self.assertEqual(parsed.follow_up_items, [])
        self.assertEqual(parsed.follow_up_questions, [])

    def test_normalize_parsed_check_in_infers_meh_sleep_as_okay(self):
        parsed = normalize_parsed_check_in(
            ParsedCheckIn(
                check_in=CheckIn(
                    raw_text=(
                        "im ok, sleep was meh, and i have 30 minutos to exercise. "
                        "I slept 5 hours, it was fine."
                    ),
                    sleep_hours=5,
                ),
                follow_up_questions=[
                    "How many hours did you sleep, and was the sleep poor, okay, or good?"
                ],
            )
        )

        self.assertEqual(parsed.check_in.sleep_quality, "okay")
        self.assertEqual(parsed.follow_up_questions, [])

    def test_follow_up_triggers_do_not_invent_body_area_questions(self):
        parsed = apply_follow_up_triggers(
            ParsedCheckIn(check_in=CheckIn(raw_text="My shoulder feels odd today."))
        )

        self.assertEqual(parsed.follow_up_questions, [])

    def test_follow_up_triggers_keep_model_activity_question_when_activity_is_present(self):
        parsed = apply_follow_up_triggers(
            ParsedCheckIn(
                check_in=CheckIn(raw_text="I played soccer yesterday."),
                follow_up_items=[
                    FollowUpQuestion(
                        field="readiness",
                        question=(
                            "How did soccer affect your energy, soreness, "
                            "and readiness today?"
                        ),
                        reason="Recent sport may affect today's training readiness.",
                    )
                ],
            )
        )

        self.assertIn(
            "How did soccer affect your energy, soreness, and readiness today?",
            parsed.follow_up_questions,
        )

    def test_follow_up_triggers_are_deduplicated(self):
        parsed = apply_follow_up_triggers(
            ParsedCheckIn(
                check_in=CheckIn(raw_text="Slept badly last night."),
                follow_up_items=[
                    FollowUpQuestion(
                        field="sleep_hours",
                        question="How many hours did you sleep?",
                        reason="Sleep duration is missing.",
                    ),
                    FollowUpQuestion(
                        field="sleep_hours",
                        question="How many hours did you sleep?",
                        reason="Duplicate from model.",
                    ),
                ],
            )
        )

        self.assertEqual(
            parsed.follow_up_questions.count("How many hours did you sleep?"),
            1,
        )

    def test_parse_model_response_rejects_invalid_json(self):
        with self.assertRaises(ValueError):
            parse_model_response("not json")

    def test_parse_model_response_logs_invalid_json_when_debug_flag_is_enabled(self):
        response_text = "not json"

        with patch.dict(
            "os.environ",
            {PARSER_LOG_RESPONSES_ENV_VAR: "1"},
            clear=True,
        ), self.assertLogs("training_coach.parser", level="ERROR") as logs:
            with self.assertRaises(ValueError):
                parse_model_response(response_text)

        output = "\n".join(logs.output)
        self.assertIn("event=parser_invalid_json_response", output)
        self.assertIn(repr(response_text), output)

    def test_parse_model_response_logs_parsed_json_when_debug_flag_is_enabled(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "60 min",
                    "time_available_minutes": 60,
                    "pain_or_injury": "unsure",
                }
            }
        )

        with patch.dict(
            "os.environ",
            {PARSER_LOG_PARSED_JSON_ENV_VAR: "1"},
            clear=True,
        ), self.assertLogs("training_coach.parser", level="INFO") as logs:
            parse_model_response(response)

        output = "\n".join(logs.output)
        self.assertIn("event=parser_json_loaded", output)
        self.assertIn("event=parser_validated_json", output)
        self.assertIn("event=parser_normalized_json", output)

    def test_negated_pain_mentions_do_not_force_injury(self):
        for text in (
            "about my calves, nothing hurts",
            "60 minutes, slept great, feeling strong, no pain",
            "shoulder doesn't hurt anymore",
            "pain free today",
        ):
            parsed = ParsedCheckIn(
                check_in=CheckIn(raw_text=text, pain_or_injury="no")
            )

            normalized = normalize_parsed_check_in(parsed)

            self.assertEqual(
                normalized.check_in.pain_or_injury, "no", msg=repr(text)
            )
            self.assertEqual(normalized.check_in.pain_issues, [], msg=repr(text))

    def test_positive_pain_mentions_still_force_injury(self):
        for text in (
            "my tricep aches a bit",
            "no pain except my knee hurts",
            "left shoulder hurts but back doesn't hurt",
        ):
            parsed = ParsedCheckIn(
                check_in=CheckIn(raw_text=text, pain_or_injury="no")
            )

            normalized = normalize_parsed_check_in(parsed)

            self.assertEqual(
                normalized.check_in.pain_or_injury, "yes", msg=repr(text)
            )

    def test_parse_model_response_drops_training_actions(self):
        response = json.dumps(
            {
                "check_in": {"raw_text": "Yesterday I ran a 10k."},
                "training_action": "skip squats",
            }
        )

        parsed = parse_model_response(response)

        # The invented key is dropped; no training decision reaches the engine.
        self.assertNotIn("training_action", parsed.model_dump())

    def test_parse_model_response_drops_synonym_keys_then_infers(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "I had a terrible night.",
                    "rest": "suboptimal",
                    "pain_or_injury": "unsure",
                }
            }
        )

        parsed = parse_model_response(response)

        # "rest" is dropped rather than rejected; deterministic inference
        # still reads "terrible" from the raw text.
        self.assertEqual(parsed.check_in.sleep_quality, "poor")

    def test_parse_model_response_coerces_synonym_values_to_null_then_infers(self):
        response = json.dumps(
            {
                "check_in": {
                    "raw_text": "I had a terrible night.",
                    "sleep_quality": "suboptimal",
                    "pain_or_injury": "unsure",
                }
            }
        )

        parsed = parse_model_response(response)

        # "suboptimal" is nulled rather than rejected, then deterministic
        # inference reads "terrible" from the raw text.
        self.assertEqual(parsed.check_in.sleep_quality, "poor")


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