tests/test_parser_runtime.py

import unittest
from unittest.mock import patch

from training_coach.parser_runtime import (
    DEFAULT_TRANSFORMERS_MODEL,
    PARSER_MODEL,
    MODEL_CACHE_ENV_VAR,
    ParserRuntimeUnavailableError,
    _load_transformers,
    generate_parser_response,
    load_parser_model,
    parse_check_in_with_model,
)


class ParserRuntimeTest(unittest.TestCase):
    def tearDown(self):
        load_parser_model.cache_clear()

    def test_runtime_exports_model_name(self):
        self.assertEqual(PARSER_MODEL, "Qwen/Qwen2.5-1.5B-Instruct")
        self.assertEqual(DEFAULT_TRANSFORMERS_MODEL, "Qwen/Qwen3-1.7B")

    def test_runtime_uses_default_huggingface_cache_by_default(self):
        self.assertEqual(MODEL_CACHE_ENV_VAR, "PARSER_MODEL_CACHE_DIR")

    def test_load_parser_model_uses_default_hf_cache_without_override(self):
        class FakeTokenizer:
            @classmethod
            def from_pretrained(cls, model_name, cache_dir=None):
                return {"model_name": model_name, "cache_dir": cache_dir}

        class FakeModel:
            @classmethod
            def from_pretrained(cls, model_name, **kwargs):
                return {"model_name": model_name, **kwargs}

        load_parser_model.cache_clear()
        with patch.dict("os.environ", {}, clear=True), patch(
            "training_coach.parser_runtime._load_transformers",
            return_value=(FakeModel, FakeTokenizer),
        ):
            tokenizer, model = load_parser_model("test/model")

        self.assertIsNone(tokenizer["cache_dir"])
        self.assertIsNone(model["cache_dir"])

    def test_load_parser_model_allows_explicit_cache_override(self):
        class FakeTokenizer:
            @classmethod
            def from_pretrained(cls, model_name, cache_dir=None):
                return {"model_name": model_name, "cache_dir": cache_dir}

        class FakeModel:
            @classmethod
            def from_pretrained(cls, model_name, **kwargs):
                return {"model_name": model_name, **kwargs}

        load_parser_model.cache_clear()
        with patch.dict(
            "os.environ",
            {MODEL_CACHE_ENV_VAR: "/tmp/parser-cache"},
            clear=True,
        ), patch(
            "training_coach.parser_runtime._load_transformers",
            return_value=(FakeModel, FakeTokenizer),
        ):
            tokenizer, model = load_parser_model("test/model")

        self.assertEqual(tokenizer["cache_dir"], "/tmp/parser-cache")
        self.assertEqual(model["cache_dir"], "/tmp/parser-cache")

    def test_runtime_can_import_transformers_or_reports_clear_error(self):
        try:
            loaded = _load_transformers()
        except ParserRuntimeUnavailableError as error:
            self.assertIn("Install transformers", str(error))
        else:
            self.assertEqual(len(loaded), 2)

    def test_runtime_functions_are_callable(self):
        self.assertTrue(callable(generate_parser_response))
        self.assertTrue(callable(parse_check_in_with_model))


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