tests/test_model_load.py

"""Tests for model loading — TRN-01 + TRN-10.

Uses sys.modules injection so tests run without unsloth installed on laptop.
"""
import sys
from unittest.mock import MagicMock, patch
import pytest
from hydra import initialize, compose


# Stub unsloth before any import of train.model_load
if "unsloth" not in sys.modules:
    sys.modules["unsloth"] = MagicMock()


def test_instruct_gate_rejects_non_instruct_variant():
    """TRN-01 gate: non-Instruct base model raises ValueError before any download."""
    with initialize(config_path="../configs", version_base="1.3"):
        cfg = compose(config_name="config", overrides=["model.name=unsloth/Qwen2.5-Coder-1.5B"])
    import importlib
    import train.model_load as m
    importlib.reload(m)
    with pytest.raises(ValueError, match="TRN-01 gate"):
        m.load_model_and_tokenizer(cfg)


def test_load_model_and_tokenizer_signature_smoke():
    """TRN-01: FastLanguageModel called with correct args when mocked."""
    with initialize(config_path="../configs", version_base="1.3"):
        cfg = compose(config_name="config")

    mock_model = MagicMock(name="model")
    mock_tok = MagicMock(name="tok")
    mock_peft_model = MagicMock(name="peft_model")

    flm_mock = sys.modules["unsloth"].FastLanguageModel
    flm_mock.from_pretrained.return_value = (mock_model, mock_tok)
    flm_mock.get_peft_model.return_value = mock_peft_model

    import importlib
    import train.model_load as m
    importlib.reload(m)

    result = m.load_model_and_tokenizer(cfg)

    assert isinstance(result, tuple) and len(result) == 2
    flm_mock.from_pretrained.assert_called_with(
        model_name="unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit",
        load_in_4bit=True,
        max_seq_length=16384,
        dtype=None,
    )


def test_smoke_config_loads_0_5b():
    """TRN-10: 0.5B smoke config resolves correctly via Hydra."""
    with initialize(config_path="../configs", version_base="1.3"):
        cfg = compose(config_name="config", overrides=["model=qwen_0_5b_smoke"])
    assert cfg.model.name == "unsloth/Qwen2.5-Coder-0.5B-Instruct-bnb-4bit"
    assert cfg.model.max_seq_length == 4096
    assert cfg.model.lora_rank == 8