tests/test_writer.py

import json

from midnight_static.schema import Genre, Script
from midnight_static.writer import (
    SMOKE_PREMISES,
    NemotronWriter,
    build_writer_prompt,
    extract_json_object,
    get_writer,
    parse_writer_json,
    script_to_json_dict,
)


def test_fixture_writer_returns_valid_scripts_for_smoke_premises():
    writer = get_writer("fixture")

    for premise in SMOKE_PREMISES:
        script = writer.write(premise, genre=Genre.WEIRD)
        script.validate()
        assert isinstance(script, Script)
        assert script.genre == Genre.WEIRD
        assert premise.split()[0].title().strip(".,:;!?") in script.title


def test_script_to_json_dict_validates_and_serializes():
    writer = get_writer("fixture")
    script = writer.write("static")

    payload = script_to_json_dict(script)

    assert payload["title"] == script.title
    assert payload["genre"] == Genre.WEIRD


def test_nemotron_writer_constructs_without_loading_model():
    writer = get_writer("nemotron")

    assert isinstance(writer, NemotronWriter)
    assert writer.model_id == "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"


def test_unknown_writer_is_rejected():
    try:
        get_writer("missing")
    except ValueError as exc:
        assert "unknown writer" in str(exc)
    else:
        raise AssertionError("expected unknown writer to fail")


def test_extract_json_object_handles_fenced_output():
    raw = '```json\n{"title": "A", "nested": {"ok": true}}\n```\nextra'

    assert extract_json_object(raw) == '{"title": "A", "nested": {"ok": true}}'


def test_parse_writer_json_round_trips_script_payload():
    script = get_writer("fixture").write("static")
    raw = "prefix " + json.dumps(script_to_json_dict(script)) + " suffix"

    parsed = parse_writer_json(raw)

    assert parsed.title == script.title
    assert parsed.genre == Genre.WEIRD


def test_writer_prompt_contains_schema_constraints():
    prompt = build_writer_prompt("static", Genre.WEIRD)

    assert "strict JSON only" in prompt
    assert '"genre": "weird"' in prompt
    assert "am_michael" in prompt
    assert "delivery" in prompt