File size: 2,638 Bytes
7aa0e08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """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()
|