tests/test_browser_e2e.py

"""Browser end-to-end test: drive the real frontend in headless Chromium
against the live gr.Server app with the tiny offline model.

Run directly: .venv/bin/python tests/test_browser_e2e.py
Screenshots land in tests/screenshots/.
"""

from __future__ import annotations

import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))
sys.path.insert(0, str(Path(__file__).parent))

from conftest import build_byte_tokenizer, build_tiny_model  # noqa: E402

PORT = 7896
BASE = f"http://127.0.0.1:{PORT}"
SHOTS = Path(__file__).parent / "screenshots"


def seg_check(locator) -> None:
    """Select a segmented-control option (its radio input is visually hidden)."""
    locator.evaluate("el => { el.checked = true; el.dispatchEvent(new Event('change')); }")


def main() -> None:
    from playwright.sync_api import sync_playwright

    from miru_tracer.core.model_manager import ModelManager

    ModelManager._model = build_tiny_model()
    ModelManager._tokenizer = build_byte_tokenizer()
    ModelManager._device = "cpu"
    ModelManager._model_name = "tiny/test-model"

    import app as jacobina_app

    jacobina_app.server.launch(
        prevent_thread_lock=True, server_name="127.0.0.1", server_port=PORT, quiet=True
    )

    SHOTS.mkdir(exist_ok=True)
    console_errors: list[str] = []

    with sync_playwright() as p:
        browser = p.chromium.launch(headless=True)
        page = browser.new_page(viewport={"width": 1500, "height": 1000})
        page.on(
            "console",
            lambda msg: console_errors.append(msg.text) if msg.type == "error" else None,
        )
        page.on("pageerror", lambda err: console_errors.append(str(err)))

        page.goto(BASE, wait_until="networkidle")
        assert "Jacobina" in page.title()

        # ---- Model view (default): status restored on load
        page.wait_for_selector(".model-pill.is-loaded", timeout=15000)
        assert "tiny/test-model" in page.text_content(".model-pill")
        assert "tiny/test-model" in page.inner_text("#view-model")
        page.screenshot(path=str(SHOTS / "01-model.png"), full_page=True)
        print("PASS model view (status restore)")

        # ---- Tokens view
        page.click("text=Tokens")
        page.fill("#view-tokens textarea", "Hello Jacobina")
        page.click("#view-tokens button:has-text('Tokenize')")
        page.wait_for_selector("#view-tokens tbody tr", timeout=10000)
        rows = page.locator("#view-tokens .candidates-scroll tbody tr").count()
        assert rows > 3, f"expected token rows, got {rows}"
        page.fill("#view-tokens input[type=number]", "72")
        page.click("#view-tokens button:has-text('Lookup token')")
        page.wait_for_selector("#view-tokens td:text-is('ID')", timeout=10000)
        page.screenshot(path=str(SHOTS / "02-tokens.png"), full_page=True)
        print(f"PASS tokens view ({rows} rows + lookup)")

        # ---- Interactive view
        page.click("text=Interactive")
        page.click("#view-interactive button:has-text('Initialize')")
        page.wait_for_selector("#view-interactive tbody tr", timeout=20000)
        candidates = page.locator("#view-interactive tbody tr").count()
        assert candidates >= 10, f"expected 10 candidates, got {candidates}"
        assert page.locator("#view-interactive tr.is-preview").count() == 1

        page.click("#view-interactive button:has-text('Next step')")
        page.wait_for_selector("#view-interactive .status-pane:has-text('Step 1 complete')", timeout=15000)
        step_shown = page.text_content("#view-interactive .pill-value")
        assert step_shown.strip() == "1", f"step display: {step_shown!r}"

        page.click("#view-interactive button:has-text('Step back')")
        page.wait_for_selector("#view-interactive .status-pane:has-text('Undone last step')", timeout=15000)

        page.click("#view-interactive button:has-text('Run')")
        page.wait_for_selector("#view-interactive .status-pane:has-text('Total steps: 10')", timeout=60000)

        # layer lens panel
        page.click("#view-interactive summary:has-text('Layer lens')")
        page.click("#view-interactive button:has-text('Refresh lens')")
        page.wait_for_selector("#view-interactive .plot-host .plotly", timeout=20000)
        page.screenshot(path=str(SHOTS / "03-interactive.png"), full_page=True)
        print("PASS interactive view (init/step/undo/run/lens plot)")

        # ---- Generate view
        page.click("text=Generate")
        page.fill("#view-generate input[type=number]", "8")  # max tokens
        page.click("#view-generate summary:has-text('Logging & visualization')")
        # uncheck stop at EOS so step count is deterministic
        page.uncheck("#view-generate .check-field:has-text('Stop at EOS') input")
        page.click("#view-generate button:has-text('Generate')")
        page.wait_for_selector("#view-generate .status-pane:has-text('Done — 8 steps')", timeout=60000)
        assert page.locator("#view-generate .plot-host .plotly").count() == 2
        continue_btn = page.locator("#view-generate button:has-text('Continue')")
        assert continue_btn.is_enabled(), "Continue should be enabled after a run"
        # editing the prompt must invalidate Continue
        page.fill("#view-generate .prompt-editor textarea", "Changed prompt")
        page.wait_for_timeout(200)
        assert not continue_btn.is_enabled(), "Continue must disable when inputs change"
        page.screenshot(path=str(SHOTS / "04-generate.png"), full_page=True)
        print("PASS generate view (stream/plots/continue-invalidation)")

        # ---- Lens view
        page.click("nav >> text=Lens")
        # add an intervention first (steer ' the', logit basis, layer 0)
        iv_panel = page.locator("#view-lens .panel-details", has=page.locator("summary:has-text('Interventions')"))
        token_input = iv_panel.locator("input[placeholder*='Paris (leading space']")
        if not token_input.is_visible():
            page.click("#view-lens summary:has-text('Interventions')")
        token_input.fill(" the")
        seg_check(iv_panel.locator(".seg-item:has-text('logit') input").first)
        page.click("#view-lens button:has-text('Add intervention')")
        page.wait_for_selector("#view-lens .iv-table-host tbody tr", timeout=10000)
        print("  intervention added")

        page.click("#view-lens button:has-text('Generate')")
        page.wait_for_selector("#view-lens .token-chip", timeout=60000)
        chips = page.locator("#view-lens .token-chip").count()
        assert chips > 10, f"expected token chips, got {chips}"
        page.wait_for_selector("#view-lens .lens-html-host table", timeout=20000)
        print(f"  generated, {chips} positions, summary table rendered")

        # click two token chips -> selection, update readouts
        page.locator("#view-lens .token-chip").nth(2).click()
        page.locator("#view-lens .token-chip").nth(3).click()
        assert "2 of" in page.text_content("#view-lens .card:has(.token-strip)")
        page.click("#view-lens button:has-text('Update readouts')")
        page.wait_for_selector(
            "#view-lens .status-pane:has-text('2 selected positions')", timeout=30000
        )
        print("  position selection + update readouts")

        # switch result views
        seg_check(page.locator("#view-lens .seg-item:has-text('Heatmap') input"))
        page.wait_for_selector(".lens-html-host p:has-text('Lens readouts')", timeout=15000)
        seg_check(page.locator("#view-lens .seg-item:has-text('Readouts') input").first)
        page.wait_for_selector(".lens-html-host p:has-text('Readout counts')", timeout=15000)
        page.screenshot(path=str(SHOTS / "05-lens.png"), full_page=True)
        print("PASS lens view (interventions/generate/selection/views)")

        # ---- Analysis view: export from interactive, re-upload
        page.click("text=Interactive")
        with page.expect_download() as dl:
            page.click("#view-interactive button:has-text('Download JSON')")
        path = dl.value.path()
        print(f"  export downloaded: {dl.value.suggested_filename}")

        page.click("text=Analysis")
        page.set_input_files("#view-analysis input[type=file]", str(path))
        page.wait_for_selector("#view-analysis .status-pane:has-text('Log loaded')", timeout=20000)
        assert page.locator("#view-analysis .plot-host .plotly").count() == 2
        assert '"schema_version": 2' in page.inner_text("#view-analysis .json-pane")
        page.screenshot(path=str(SHOTS / "06-analysis.png"), full_page=True)
        print("PASS analysis view (export round-trip)")

        # ---- theme toggle
        before = page.get_attribute("html", "data-theme")
        page.click(".icon-btn[title='Toggle theme']")
        after = page.get_attribute("html", "data-theme")
        assert {before, after} == {"dark", "light"}, (before, after)
        page.screenshot(path=str(SHOTS / "07-other-theme.png"), full_page=True)
        print("PASS theme toggle")

        browser.close()

    benign = re.compile(r"favicon|manifest|Failed to load resource.*404")
    real_errors = [e for e in console_errors if not benign.search(e)]
    assert not real_errors, f"console errors: {real_errors}"
    print("PASS no console errors")
    print("\nALL BROWSER E2E TESTS PASS")


def test_browser_e2e():
    main()


if __name__ == "__main__":
    main()