tests/test_trace.py

"""Agent-trace export + redaction (server/trace.py) and the UI download handler."""
import json

from server import events as bus
from server import trace
from ui.blocks import _on_analyze, _on_download_trace


def _drive_a_run():
    # stub mode is forced by conftest -> this produces a real run on the bus
    list(_on_analyze("Mom: dinner Sunday 6pm", None, None))


def test_export_run_shape_after_analyze():
    _drive_a_run()
    t = trace.export_run()
    assert t["schema"] == "imessage-cal-trace"
    assert t["version"] == 1
    assert t["run_id"]
    assert t["steps"], "expected at least one step from the run"
    assert t["summary"]["events"] >= 1
    allowed = {
        "stage", "level", "ts", "latency_ms", "events",
        "conflicts", "images", "tokens", "message",
    }
    for step in t["steps"]:
        assert set(step).issubset(allowed), f"unexpected keys: {set(step) - allowed}"


def test_redaction_strips_chat_names():
    with bus.run_scope("analyze"):
        bus.emit("ingest", "2 msg(s) from 3rd grade chat", images=0)
        bus.emit("decision", "1 event(s) detected", events=1)
    rid = trace.export_run()["run_id"]

    redacted = trace.export_run(run_id=rid, redact=True)
    assert all("3rd grade chat" not in s["message"] for s in redacted["steps"])
    assert redacted["redacted"] is True

    full = trace.export_run(run_id=rid, redact=False)
    assert any("3rd grade chat" in s["message"] for s in full["steps"])


def test_empty_envelope_when_no_run():
    # a run id that doesn't exist -> valid but empty
    t = trace.export_run(run_id="nonexistent")
    assert t["steps"] == []
    assert t["summary"]["events"] == 0


def test_write_trace_roundtrips(tmp_path):
    _drive_a_run()
    t = trace.export_run()
    path = trace.write_trace(t, path=str(tmp_path / "t.json"))
    with open(path, encoding="utf-8") as f:
        assert json.load(f) == t


def test_download_handler_paths():
    _drive_a_run()
    path, msg = _on_download_trace(True)
    assert path and "step" in msg

    bus._BUF.clear()  # simulate a fresh process: no runs to export
    none_path, none_msg = _on_download_trace(True)
    assert none_path is None and "No agent run" in none_msg