tests/test_render.py

"""The renderers, over every recorded state, and against text a visitor could type."""

from __future__ import annotations

import copy
import json
import re
from pathlib import Path

import pytest

import recorded
from steering.core import ALPHAS
from viz import render as R

TRACES = [t for rec in recorded.records().values() for t in rec["scenarios"]]


def panels(trace, alpha):
    return [R.formula_md(trace, alpha), R.cards_html(trace, alpha), R.response_html(trace, alpha),
            R.chart_html(trace, alpha), R.sweep_html(trace, alpha), R.schematic_html(trace, alpha),
            R.takeaway_html(trace), R.contrast_html(trace), R.explain_html(trace)]


@pytest.mark.parametrize("trace", TRACES, ids=[t["model"].split("/")[-1] + ":" + t["id"] for t in TRACES])
def test_every_state_renders_and_shows_its_own_numbers(trace):
    for alpha in ALPHAS:
        state = R.state_at(trace, alpha)
        out = panels(trace, alpha)
        token = R.label(trace["candidates"][state["selected_index"]])
        assert token in out[1] and token in out[2]
        assert R.fmt_pct(state["percents"][state["selected_index"]]) in out[3]
        # Each bar is as wide as its probability, and its marker sits at the alpha = 0 value.
        widths = [float(w) for w in re.findall(r'sd-fill [a-z]+" style="width:([0-9.]+)%', out[3])]
        assert widths == pytest.approx([*state["percents"], state["other"]], abs=1e-3)
        assert f"α = {R.fmt_alpha(alpha)}" in out[1]


def test_html_panels_carry_no_template_syntax():
    """In Gradio 6 gr.HTML is a template: a stray ${...} or {{...}} would be evaluated."""
    for trace in TRACES:
        for alpha in (-2.0, 0.0, 2.0):
            for html in panels(trace, alpha)[1:]:
                assert "{" not in html and "}" not in html and "$" not in html


def test_text_a_visitor_typed_is_escaped():
    trace = copy.deepcopy(TRACES[0])
    hostile = '<script>alert(1)</script> ${window.x} {{y}}'
    trace.update(prompt=hostile, prefix=hostile, custom=True, source="live", cached=False, stale=False)
    for alpha in (0.0, 2.0):
        for html in panels(trace, alpha)[1:]:
            assert "<script>" not in html and "${" not in html and "{{" not in html


def test_a_token_without_a_leading_space_follows_the_opening_words_directly():
    trace = copy.deepcopy(TRACES[0])
    trace["candidates"][0] = ","
    for s in trace["states"]:
        if s["selected_index"] == 0:
            s["continuation"] = ", and then" + s["continuation"][len("terrible"):]
    lead, head, _ = R.pieces(trace, R.state_at(trace, -2.0))
    assert lead == trace["prefix"] and head == ","


def test_a_live_result_never_shows_the_recorded_conclusion():
    trace = copy.deepcopy(TRACES[0])
    authored = trace["takeaway"]
    trace.update(source="live", cached=False, stale=False, takeaway="The Movie Critic example, measured live.")
    html = R.takeaway_html(trace)
    assert authored not in html and "Measured live" in html and "not repeated" in html
    trace.update(custom=True, takeaway="Your prompt, steered along the Movie Critic direction.")
    assert "not repeated" not in R.takeaway_html(trace)  # a visitor's own text has no recorded summary


@pytest.mark.parametrize("flags,expected", [
    ({"source": "recorded"}, "Recorded"),
    ({"source": "live", "cached": False, "stale": False}, "Live"),
    ({"source": "live", "cached": True, "stale": False, "age_seconds": 600}, "Replayed"),
    ({"source": "live", "cached": True, "stale": True, "age_seconds": 600}, "Fallback"),
])
def test_the_source_card_never_claims_more_than_happened(flags, expected):
    trace = {**TRACES[0], **flags}
    value, key, caption = R.source_card(trace)
    assert value == expected
    if expected == "Replayed":
        assert "10 min ago" in caption


def test_the_status_line_names_what_changed_and_nothing_otherwise():
    trace = TRACES[0]
    same = R.status_html(trace, trace["model"], trace["id"], trace["prompt"], "  " + trace["prefix"] + " ")
    assert same == ""
    changed = R.status_html(trace, trace["model"], "storyteller", "Another prompt.", trace["prefix"])
    assert "the scenario and the text" in changed


def test_padded_rows_are_declared():
    trace = copy.deepcopy(TRACES[0])
    for s in trace["states"]:
        s["selected_index"] = 1
    assert "Fewer than three tokens win" in R.chart_html(trace, 0.0)
    assert "Fewer than three tokens win" not in R.chart_html(TRACES[0], 0.0)


def test_pole_colours_follow_the_winners_at_each_end_not_the_row_order():
    """A padded row never takes a pole's colour; the winner at +2 is the positive pole wherever its row is."""
    trace = copy.deepcopy(TRACES[0])
    trace["candidates"] = [" terrible", " delicious", " awful"]   # "awful" padded in, never wins
    for s in trace["states"]:
        s["selected_index"] = 0 if s["alpha"] < 0 else 1
    assert R.poles(trace) == ["neg", "pos", "mid"]
    for s in trace["states"]:   # a direction that never changes the winner has no positive pole
        s["selected_index"] = 1
    assert R.poles(trace) == ["mid", "neg", "mid"]


def test_whitespace_tokens_are_visible_and_distinct():
    labels = [R.label(t) for t in (" \n\n", " ", " \n", "\t", " cat")]
    assert labels == ["↵↵", "␣", "↵", "⇥", "cat"]
    assert len(set(labels)) == len(labels)


def test_the_highlight_covers_the_token_and_nothing_after_it():
    trace = copy.deepcopy(TRACES[0])
    trace["candidates"][0] = ")\n\n\n"
    state = R.state_at(trace, -2.0)
    state["selected_index"] = 0
    state["continuation"] = ") Then the credits rolled."
    lead, head, tail = R.pieces(trace, state)
    assert head == ")" and tail == " Then the credits rolled."


def test_trace_files_are_named_by_content():
    first, again = R.trace_file(TRACES[0]), R.trace_file(TRACES[0])
    assert first == again and first != R.trace_file(TRACES[1])


def test_sessions_writing_the_same_trace_at_once_do_not_collide():
    import threading
    from concurrent.futures import ThreadPoolExecutor

    for round_ in range(20):
        trace = {**TRACES[0], "measured": f"written by eight sessions at once, round {round_}"}
        start = threading.Barrier(8)

        def write(_):
            start.wait()
            return R.trace_file(trace)

        with ThreadPoolExecutor(8) as pool:
            paths = set(pool.map(write, range(8)))
        assert len(paths) == 1 and json.loads(Path(paths.pop()).read_text(encoding="utf-8")) == trace
    assert not list(R.TRACE_DIR.glob("*.tmp"))


def test_old_trace_files_are_removed():
    import os
    import time

    old = Path(R.trace_file({**TRACES[0], "measured": "an old trace"}))
    os.utime(old, (time.time() - 2 * R.TRACE_MAX_AGE_SECONDS,) * 2)
    R.trace_file(TRACES[1])
    assert not old.exists()


@pytest.mark.parametrize("token,continuation", [("\n\n", "The end."), ("<|im_end|>", "")])
def test_a_first_token_the_response_does_not_contain_is_shown_by_its_label(token, continuation):
    # A newline is collapsed out of the response text and an end-of-text token is skipped in decoding.
    trace = copy.deepcopy(TRACES[0])
    state = R.state_at(trace, -2.0)
    trace["candidates"][state["selected_index"]] = token
    state["continuation"] = continuation
    lead, head, tail = R.pieces(trace, state)
    assert (lead, head) == (trace["prefix"], R.label(token)) and tail.strip() == continuation
    assert "␣" not in R.response_html(trace, -2.0)