tests/test_streaming.py

"""Unit tests for the streaming reflection path.

These tests cover the pure-Python pieces of the streaming refactor:
section flushing, sanitizer ordering, and safety short-circuit. They
do not touch FastAPI, gradio.Server, or any LLM backend — the
streaming dispatcher is mocked.
"""
import pytest

from app import _section_flush_loop


def _tokens(*chunks):
    """Helper: turn a list of strings into an iterator."""
    for c in chunks:
        yield c


def test_section_flush_emits_one_event_per_heading():
    stream = _tokens(
        "## Mirror\n", "One possible ", "reading is...\n",
        "## Key Symbols\n", "- forest\n",
    )
    events = list(_section_flush_loop(stream))
    headings = [e["heading"] for e in events]
    assert headings == ["Mirror", "Key Symbols"]


def test_section_flush_runs_sanitizer_per_section():
    # "You should" at sentence-start in Mirror; rewrite must apply
    # when Mirror flushes (i.e. when ## Key Symbols arrives).
    stream = _tokens(
        "## Mirror\n", "You should ", "notice the forest.\n",
        "## Key Symbols\n", "- forest\n",
    )
    events = list(_section_flush_loop(stream))
    mirror = next(e for e in events if e["heading"] == "Mirror")
    assert "you should" not in mirror["markdown"].lower()
    assert "you might notice" in mirror["markdown"].lower()


def test_section_flush_emits_trailing_section():
    # The final section has no following ## boundary — must still flush.
    stream = _tokens(
        "## Mirror\n", "...\n",
        "## Gentle Question\n", "What might you notice?\n",
    )
    events = list(_section_flush_loop(stream))
    assert events[-1]["heading"] == "Gentle Question"
    assert "What might you notice?" in events[-1]["markdown"]


def test_section_flush_handles_split_heading():
    # The `## ` boundary may straddle two tokens.
    stream = _tokens("## Mirror\n...\n#", "# Key Symbols\n- forest\n")
    events = list(_section_flush_loop(stream))
    assert [e["heading"] for e in events] == ["Mirror", "Key Symbols"]


def test_section_flush_skips_preamble_before_first_heading():
    # The model sometimes emits a stray newline before "## Mirror".
    # Anything before the first ## is dropped.
    stream = _tokens("\n\n## Mirror\nbody\n## Symbols\n- x\n")
    events = list(_section_flush_loop(stream))
    assert [e["heading"] for e in events] == ["Mirror", "Symbols"]


def test_section_flush_does_not_match_h3_subheading():
    """### Subheading inside a section body must not be parsed as a new section.

    Regression: text.find("## ") matched the inner `## ` inside `###`.
    """
    stream = _tokens(
        "## Mirror\n",
        "Some preamble.\n",
        "### subsection\n",
        "body inside subsection\n",
        "## Symbols\n",
        "- forest\n",
    )
    events = list(_section_flush_loop(stream))
    headings = [e["heading"] for e in events]
    assert headings == ["Mirror", "Symbols"]
    mirror = next(e for e in events if e["heading"] == "Mirror")
    assert "### subsection" in mirror["markdown"]
    assert "body inside subsection" in mirror["markdown"]


# ---------------------------------------------------------------------------
# _run_llama_cpp_stream tests (Task 2)
# ---------------------------------------------------------------------------
from unittest.mock import MagicMock, patch


class _FakeLlama:
    """Stand-in for llama_cpp.Llama with create_chat_completion(stream=True)."""
    def __init__(self, chunks):
        self._chunks = chunks
        self.last_kwargs = None

    def create_chat_completion(self, **kwargs):
        self.last_kwargs = kwargs
        for piece in self._chunks:
            yield {"choices": [{"delta": {"content": piece}}]}


def test_run_llama_cpp_stream_yields_token_chunks():
    from app import _run_llama_cpp_stream

    fake = _FakeLlama(["## Mirror\n", "hello", " world\n"])
    with patch("app._load_llama_cpp_model", return_value=(fake, None)):
        out = list(_run_llama_cpp_stream(
            text="I dreamt of a forest",
            entry_type="Dream",
            depth=2,
            symbol_matches=[],
            grounded_jungian=False,
            include_question=True,
        ))

    assert out == ["## Mirror\n", "hello", " world\n"]
    assert fake.last_kwargs["stream"] is True


def test_run_llama_cpp_stream_surfaces_load_error():
    from app import _run_llama_cpp_stream

    with patch("app._load_llama_cpp_model", return_value=(None, "boom")):
        out = list(_run_llama_cpp_stream(
            text="x", entry_type="Dream", depth=2,
            symbol_matches=[], grounded_jungian=False, include_question=True,
        ))

    # On load error, the streamer emits a single error sentinel and stops.
    assert len(out) == 1
    assert out[0].startswith("__error__:")
    assert "boom" in out[0]


# ---------------------------------------------------------------------------
# run_model_stream dispatcher tests (Task 3)
# ---------------------------------------------------------------------------

def test_run_model_stream_dispatches_to_llama_cpp_by_default(monkeypatch):
    import app as appmod

    monkeypatch.setattr(appmod, "BACKEND", "llama_cpp")
    called = []

    def fake_stream(*args, **kwargs):
        called.append(("llama_cpp", args, kwargs))
        yield "hello"

    monkeypatch.setattr(appmod, "_run_llama_cpp_stream", fake_stream)

    out = list(appmod.run_model_stream(
        text="x", entry_type="Dream", depth=2,
        symbol_matches=[], grounded_jungian=False, include_question=True,
    ))
    assert out == ["hello"]
    assert called[0][0] == "llama_cpp"


def test_run_model_stream_dispatches_to_ollama_when_configured(monkeypatch):
    import app as appmod

    monkeypatch.setattr(appmod, "BACKEND", "ollama")
    called = []

    def fake_stream(*args, **kwargs):
        called.append(("ollama",))
        yield "world"

    monkeypatch.setattr(appmod, "_run_ollama_stream", fake_stream)

    out = list(appmod.run_model_stream(
        text="x", entry_type="Dream", depth=2,
        symbol_matches=[], grounded_jungian=False, include_question=True,
    ))
    assert out == ["world"]
    assert called[0][0] == "ollama"


# ---------------------------------------------------------------------------
# reflect_api generator tests (Task 4)
# ---------------------------------------------------------------------------

def test_reflect_api_safety_short_circuits_before_llm(monkeypatch):
    import app as appmod

    llm_called = []
    def fake_stream(*a, **kw):
        llm_called.append(True)
        yield "## Mirror\nshould never appear\n"

    monkeypatch.setattr(appmod, "run_model_stream", fake_stream)

    events = list(appmod.reflect_api(
        entry="I want to end it all tonight.",
        entry_type="Journal",
        depth=2,
        grounded_jungian=False,
        include_question=True,
    ))
    types = [e["event"] for e in events]
    assert "safety" in types
    assert "reading_section" not in types
    assert llm_called == []  # LLM must not have been called


def test_reflect_api_emits_expected_event_sequence(monkeypatch):
    import app as appmod

    def fake_stream(*a, **kw):
        yield "## Mirror\nA possible reading.\n"
        yield "## Key Symbols\n- forest\n"

    monkeypatch.setattr(appmod, "run_model_stream", fake_stream)
    monkeypatch.setattr(appmod, "render_mandala_to_path", lambda s: "fake.png")

    events = list(appmod.reflect_api(
        entry="I dreamt of a forest with a river.",
        entry_type="Dream",
        depth=2,
        grounded_jungian=False,
        include_question=True,
    ))
    types = [e["event"] for e in events]
    # Order: symbols first, then mandala, then per-section, then done.
    assert types[0] == "symbols"
    assert types[1] == "mandala"
    assert types.count("reading_section") == 2
    assert types[-1] == "done"


def test_reflect_api_surfaces_llm_error_sentinel(monkeypatch):
    import app as appmod

    def fake_stream(*a, **kw):
        yield "__error__:llama.cpp call failed (RuntimeError: oom)."

    monkeypatch.setattr(appmod, "run_model_stream", fake_stream)
    monkeypatch.setattr(appmod, "render_mandala_to_path", lambda s: "fake.png")

    events = list(appmod.reflect_api(
        entry="x", entry_type="Journal", depth=2,
        grounded_jungian=False, include_question=True,
    ))
    types = [e["event"] for e in events]
    assert "error" in types
    err = next(e for e in events if e["event"] == "error")
    assert "llama.cpp call failed" in err["message"]


def test_reflect_api_rejects_empty_entry():
    import app as appmod
    events = list(appmod.reflect_api(
        entry="   ", entry_type="Journal", depth=2,
        grounded_jungian=False, include_question=True,
    ))
    assert events == [{"event": "error", "code": "empty",
                       "message": "Write something first."}]


# ---------------------------------------------------------------------------
# Task 4b: MAX_SECTION_BUFFER safety valve + reflect_api coverage gaps
# ---------------------------------------------------------------------------

def test_section_flush_max_buffer_emits_synthetic_heading():
    """If no ## boundary arrives, flush at MAX_SECTION_BUFFER as Stream."""
    from app import _section_flush_loop, MAX_SECTION_BUFFER

    big_blob = "x" * (MAX_SECTION_BUFFER + 100)
    stream = _tokens("## Mirror\n", big_blob)
    events = list(_section_flush_loop(stream))
    headings = [e["heading"] for e in events]
    # Mirror flushes when the synthetic boundary trips, with heading kept.
    assert "Mirror" in headings
    # And no event's body exceeds MAX_SECTION_BUFFER materially.
    assert all(len(e["markdown"]) <= MAX_SECTION_BUFFER * 2 for e in events)


def test_reflect_api_continues_after_mandala_failure(monkeypatch):
    """A mandala render failure emits an error event but does NOT short-circuit.

    The reading is the primary output; the mandala is decoration. Subsequent
    reading_section events must still flow.
    """
    import app as appmod

    def fake_stream(*a, **kw):
        yield "## Mirror\nA possible reading.\n"

    def boom(symbol_matches):
        raise RuntimeError("PIL exploded")

    monkeypatch.setattr(appmod, "run_model_stream", fake_stream)
    monkeypatch.setattr(appmod, "render_mandala_to_path", boom)

    events = list(appmod.reflect_api(
        entry="I dreamt of a forest with a river.",
        entry_type="Dream",
        depth=2,
        grounded_jungian=False,
        include_question=True,
    ))
    types = [e["event"] for e in events]
    # Symbols first, then error (from mandala), then reading_section, then done.
    assert "symbols" in types
    err = next(e for e in events if e["event"] == "error")
    assert err["code"] == "mandala_failed"
    assert "PIL exploded" in err["message"]
    assert "reading_section" in types  # Reading still flows.
    assert types[-1] == "done"


def test_reflect_api_handles_empty_stream(monkeypatch):
    """If run_model_stream yields nothing, emit empty_stream error + done."""
    import app as appmod

    def fake_stream(*a, **kw):
        if False:
            yield  # empty generator
        return

    monkeypatch.setattr(appmod, "run_model_stream", fake_stream)
    monkeypatch.setattr(appmod, "render_mandala_to_path", lambda s: "fake.png")

    events = list(appmod.reflect_api(
        entry="I dreamt of a forest with a river.",
        entry_type="Dream",
        depth=2,
        grounded_jungian=False,
        include_question=True,
    ))
    types = [e["event"] for e in events]
    err = next(e for e in events if e["event"] == "error")
    assert err["code"] == "empty_stream"
    assert types[-1] == "done"
    assert "reading_section" not in types