tests/test_sft_scaffold.py

"""SFT scaffold tests — TRN-02.

Uses sys.modules injection so tests run without trl/unsloth/peft installed.
All tests complete in <3s on CPU (no GPU required).
"""
import sys
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from hydra import initialize, compose


# ── Stub GPU-only modules before any import of train.sft ────────────────────

def _stub_heavy_deps():
    """Inject stubs for trl, datasets, unsloth, peft so tests work on CPU."""
    for mod in ["trl", "datasets", "unsloth", "peft", "huggingface_hub"]:
        if mod not in sys.modules:
            sys.modules[mod] = MagicMock()


_stub_heavy_deps()


def _build_cfg(tmp_path, hub_push=False):
    with initialize(config_path="../configs", version_base="1.3"):
        cfg = compose(
            config_name="config",
            overrides=[
                "train=sft",
                f"output_dir={tmp_path}",
                f"data.sft_traces_path={tmp_path}/fake_sft.jsonl",
                f"+hub.push={str(hub_push).lower()}",
                "+hub.repo_id=test/fathom-sft",
            ],
        )
    return cfg


def _make_fake_dataset():
    return [{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}]


def _make_mocks():
    model_mock = MagicMock()
    tok_mock = MagicMock()
    tok_mock.apply_chat_template.return_value = (
        "<|im_start|>user\nhi<|im_end|>\n<|im_start|>assistant\nhello<|im_end|>"
    )
    return model_mock, tok_mock


def test_formatting_func_uses_chat_template(tmp_path):
    """Test A: formatting_func uses apply_chat_template (STACK §3.3)."""
    cfg = _build_cfg(tmp_path)
    model_mock, tok_mock = _make_mocks()
    dataset = _make_fake_dataset()

    # Stub trainer to capture formatting_func
    trainer_mock = MagicMock()
    trainer_mock.train.return_value = None
    captured_fmt_fn = {}

    def fake_sft_trainer(*args, **kwargs):
        captured_fmt_fn["fn"] = kwargs.get("formatting_func")
        return trainer_mock

    sys.modules["trl"].SFTTrainer = MagicMock(side_effect=fake_sft_trainer)
    sys.modules["trl"].SFTConfig = MagicMock(return_value=MagicMock())

    # Reload to pick up stubs
    import importlib
    import train.sft as sft_module
    importlib.reload(sft_module)

    sft_module.run_sft(cfg, model_mock, tok_mock, dataset=dataset)

    assert "fn" in captured_fmt_fn
    fmt_fn = captured_fmt_fn["fn"]
    assert callable(fmt_fn)
    result = fmt_fn(dataset[0])
    assert "<|im_start|>" in result
    assert "<|im_end|>" in result
    tok_mock.apply_chat_template.assert_called()
    call_kwargs = tok_mock.apply_chat_template.call_args.kwargs
    assert call_kwargs.get("tokenize") is False
    assert call_kwargs.get("add_generation_prompt") is False


def test_sft_config_mirrors_hydra_yaml(tmp_path):
    """Test B: SFTConfig kwargs mirror configs/train/sft.yaml."""
    cfg = _build_cfg(tmp_path)
    model_mock, tok_mock = _make_mocks()
    dataset = _make_fake_dataset()

    captured_sft_config_kwargs = {}

    def fake_sft_config(**kwargs):
        captured_sft_config_kwargs.update(kwargs)
        return MagicMock()

    trainer_mock = MagicMock()
    trainer_mock.train.return_value = None
    sys.modules["trl"].SFTTrainer = MagicMock(return_value=trainer_mock)
    sys.modules["trl"].SFTConfig = MagicMock(side_effect=fake_sft_config)

    import importlib
    import train.sft as sft_module
    importlib.reload(sft_module)

    sft_module.run_sft(cfg, model_mock, tok_mock, dataset=dataset)

    assert captured_sft_config_kwargs.get("learning_rate") == pytest.approx(2.0e-4)
    assert captured_sft_config_kwargs.get("num_train_epochs") == 1
    assert captured_sft_config_kwargs.get("max_seq_length") == 8192
    assert captured_sft_config_kwargs.get("per_device_train_batch_size") == 2
    assert captured_sft_config_kwargs.get("gradient_accumulation_steps") == 4
    assert captured_sft_config_kwargs.get("optim") == "adamw_8bit"
    assert captured_sft_config_kwargs.get("bf16") is True
    assert captured_sft_config_kwargs.get("seed") == 42


def test_adapter_save_always_happens_hub_push_is_gated(tmp_path, monkeypatch):
    """Test C: adapter save unconditional; Hub push gated on hub.push + HF_TOKEN."""
    cfg = _build_cfg(tmp_path, hub_push=True)
    model_mock, tok_mock = _make_mocks()
    dataset = _make_fake_dataset()

    monkeypatch.delenv("HF_TOKEN", raising=False)

    trainer_mock = MagicMock()
    trainer_mock.train.return_value = None
    sys.modules["trl"].SFTTrainer = MagicMock(return_value=trainer_mock)
    sys.modules["trl"].SFTConfig = MagicMock(return_value=MagicMock())

    import importlib
    import train.sft as sft_module
    importlib.reload(sft_module)

    result_path = sft_module.run_sft(cfg, model_mock, tok_mock, dataset=dataset)

    assert model_mock.save_pretrained.call_count >= 1
    saved_path = str(model_mock.save_pretrained.call_args_list[0][0][0])
    assert "sft_adapter" in saved_path
    assert model_mock.push_to_hub.call_count == 0
    assert str(result_path).endswith("sft_adapter")