"""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()