tests/test_parser_ollama.py

import json
import unittest
from unittest.mock import patch

from training_coach.parser_ollama import (
    OLLAMA_MODEL,
    OLLAMA_TINY_CANDIDATE,
    build_ollama_prompt,
    generate_parser_response_ollama,
    parse_check_in_with_ollama,
)


class _FakeResponse:
    def __init__(self, payload):
        self.payload = payload

    def __enter__(self):
        return self

    def __exit__(self, *args):
        return False

    def read(self):
        return json.dumps(self.payload).encode("utf-8")


class OllamaParserTest(unittest.TestCase):
    def test_ollama_model_choices_are_locked(self):
        self.assertEqual(OLLAMA_MODEL, "qwen3:1.7B")
        self.assertEqual(OLLAMA_TINY_CANDIDATE, "qwen3:4b")

    def test_build_ollama_prompt_includes_schema_and_check_in(self):
        prompt = build_ollama_prompt("Yesterday I ran a 10k.")

        self.assertIn("Expected response JSON schema", prompt)
        self.assertIn("pain_issues", prompt)
        self.assertIn("Yesterday I ran a 10k.", prompt)

    def test_generate_parser_response_ollama_sends_schema_format(self):
        fake_payload = {
            "prompt_eval_count": 120,
            "eval_count": 40,
            "total_duration": 1_500_000_000,
            "load_duration": 100_000_000,
            "prompt_eval_duration": 400_000_000,
            "eval_duration": 900_000_000,
            "response": json.dumps(
                {
                    "check_in": {
                        "raw_text": "60 min",
                        "time_available_minutes": 60,
                    }
                }
            )
        }

        with patch(
            "training_coach.parser_ollama.request.urlopen",
            return_value=_FakeResponse(fake_payload),
        ) as urlopen, self.assertLogs(
            "training_coach.parser_ollama",
            level="INFO",
        ) as logs:
            response = generate_parser_response_ollama("60 min")

        sent_request = urlopen.call_args.args[0]
        body = json.loads(sent_request.data.decode("utf-8"))
        self.assertEqual(body["model"], "qwen3:1.7B")
        self.assertEqual(body["stream"], False)
        self.assertEqual(body["think"], False)
        self.assertEqual(body["format"]["title"], "ParsedCheckIn")
        self.assertIn("60 min", body["prompt"])
        self.assertIn("check_in", response)
        output = "\n".join(logs.output)
        self.assertIn("event=parser_ollama_request", output)
        self.assertIn("event=parser_ollama_response", output)
        self.assertIn("prompt_eval_count=120", output)
        self.assertIn("eval_count=40", output)

    def test_generate_parser_response_ollama_logs_private_trace_when_enabled(self):
        fake_payload = {"response": '{"check_in": {"raw_text": "60 min"}}'}

        with patch.dict(
            "os.environ",
            {
                "PARSER_LOG_MESSAGES": "1",
                "PARSER_LOG_RESPONSES": "1",
            },
            clear=True,
        ), patch(
            "training_coach.parser_ollama.request.urlopen",
            return_value=_FakeResponse(fake_payload),
        ), self.assertLogs("training_coach.parser", level="INFO") as logs:
            generate_parser_response_ollama("60 min")

        output = "\n".join(logs.output)
        self.assertIn("event=parser_message", output)
        self.assertIn("event=parser_response_text", output)
        self.assertIn("60 min", output)

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

        with patch(
            "training_coach.parser_ollama.request.urlopen",
            return_value=_FakeResponse(fake_payload),
        ):
            parsed = parse_check_in_with_ollama("60 min")

        self.assertEqual(parsed.check_in.time_available_minutes, 60)


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