src/understudy_test.py

"""Tests for the Understudy ChatML prompt format + completion plumbing.
OTW_SKIP_PREFETCH avoids the import-time GGUF download; the llama.cpp call
is faked, so this runs with no model and no llama_cpp installed.

Run:  .venv/bin/python src/understudy_test.py   (or pytest)
"""

import os
import sys

os.environ["OTW_SKIP_PREFETCH"] = "1"
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src import understudy  # noqa: E402
from src.prompts import UNDERSTUDY_SYSTEM  # noqa: E402

RAW = """GENRE: Drill
TITLE: Tiny But Mighty
LYRICS:
[verse]
half a billion params, still I spit,
your job description? perfect fit.
[chorus]
open to work, open to win,
let the understudy in."""


def test_messages_carry_system_and_payload():
    msgs = understudy._messages("MY RESUME", "MY JOB", "🎤 Drill", 8, "🔥 unhinged — x")
    assert msgs[0] == {"role": "system", "content": UNDERSTUDY_SYSTEM}
    u = msgs[1]
    assert u["role"] == "user"
    assert "MY RESUME" in u["content"] and "MY JOB" in u["content"]
    assert "SENDABILITY: 8/10" in u["content"]  # the build_messages user payload


def test_write_uses_chat_completion_and_strips():
    calls = {}

    class FakeLlama:
        def create_chat_completion(self, **kw):
            calls.update(kw)
            return {"choices": [{"message": {"content": RAW + "\n"}}]}

    understudy._get_llm = lambda: FakeLlama()
    out = understudy.write("r", "j", "🎤 Drill", 8, "🔥 unhinged — x")
    assert out == RAW  # stripped
    assert calls["stop"] == understudy._STOP
    assert calls["temperature"] == understudy._TEMPERATURE
    assert calls["repeat_penalty"] == understudy._REPEAT_PENALTY
    # passed as chat messages (so llama.cpp applies the ChatML template + specials)
    assert calls["messages"][0]["role"] == "system"
    assert "<|im_start|>" not in str(calls["messages"])  # never hand-rolled into the text


def test_prefetch_skipped_sets_load_error():
    # imported with OTW_SKIP_PREFETCH=1 -> graceful unavailable, no crash
    assert understudy.load_error is not None
    assert "prefetch skipped" in str(understudy.load_error)


def _run():
    fns = {k: v for k, v in globals().items() if k.startswith("test_") and callable(v)}
    failed = 0
    for name, fn in fns.items():
        try:
            fn()
            print(f"PASS {name}")
        except Exception as e:
            failed += 1
            import traceback
            print(f"FAIL {name}: {e}")
            traceback.print_exc()
    print(f"\n{len(fns) - failed}/{len(fns)} passed")
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    _run()