stack/tests/notebooks/test_code_families_nb.py

"""notebooks/code_families.py — structure and headless degradation."""

from __future__ import annotations

import ast
import importlib.util
import re
from pathlib import Path

import marimo as mo
import pytest

NB = Path(__file__).resolve().parents[2] / "notebooks" / "code_families.py"
_ROOT = Path(__file__).resolve().parents[2]
#: The two live tests run the whole notebook against the real replica and
#: bibliography (data symlinks in a checkout); CI has neither, so they skip
#: there — the headless-degrade test above covers the no-data path.
_HAS_DATA = (_ROOT / "data" / "replica" / "aco.ro.duckdb").exists() and (
    _ROOT / "data" / "bib.sqlite"
).exists()
_live = pytest.mark.skipif(not _HAS_DATA, reason="needs the live replica and bib")


def _load():
    spec = importlib.util.spec_from_file_location("code_families_nb", NB)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def test_notebook_is_a_marimo_app():
    mod = _load()
    assert mod.app.__class__.__name__ == "App"


def test_cells_are_anonymous_and_banners_present():
    src = NB.read_text()
    tree = ast.parse(src)
    names = [n.name for n in tree.body if isinstance(n, ast.FunctionDef)]
    assert names and set(names) == {"_"}
    banners = re.findall(r"# ── (\S+)\. ", src)
    assert banners == [
        "0",
        "1",
        "2",
        "3",
        "4",
        "5",
        "6",
        "7",
        "7b",
        "7c",
        "7d",
        "7e",
        "7f",
        "8",
    ]


def test_headless_run_degrades_without_data(monkeypatch, tmp_path):
    """With no replica and no bib the guards render notes instead of raising."""
    monkeypatch.setenv("STACK_DUCKDB_REPLICA", "1")
    mod = _load()
    import conf.connect as cc

    monkeypatch.setattr(
        cc,
        "duckdb",
        lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError("no replica")),
    )
    monkeypatch.setattr(
        cc, "bib", lambda *a, **k: (_ for _ in ()).throw(FileNotFoundError("no bib"))
    )
    result = mod.app.run()  # marimo runs all cells; mo.stop cascades are fine
    assert result is not None


def _tables_in(obj):
    """Every mo.ui.table nested under a marimo cell output — mo.vstack
    (and other containers) expose their children as live objects via
    ``_live_children``, not just pre-rendered HTML, so a table's full
    ``.data`` (not the ~10-row preview ``_repr_html_`` embeds) is
    reachable for assertions."""
    found = []
    if isinstance(obj, mo.ui.table):
        found.append(obj)
    for child in getattr(obj, "_live_children", None) or []:
        found.extend(_tables_in(child))
    return found


@_live
def test_section5_tolerates_missing_llm_db_password(monkeypatch):
    """Section 5's pgvector chunk-count lookup must degrade to a note,
    never raise, when LLM_DB_PASSWORD is unset — the bib item-tag counts
    still render (task-8-context.md: "the section-5 cell must render the
    bib tag counts with a note when it is missing, and never raise").

    The note-text assertions alone are vacuous — they'd pass even if
    ``store.list_tags`` raised, since the note text sits in the cell's
    intro prose. So also require the ``family:*`` tag table to have
    actually rendered (a mo.ui.table only appears on the list_tags
    success path — the ``except`` branch builds no table at all) and to
    carry a real, live row: ``family:CCM`` with a positive count.
    """
    monkeypatch.delenv("LLM_DB_PASSWORD", raising=False)
    mod = _load()
    outputs, _defs = mod.app.run()  # real replica/bib — data symlinks are in place
    rendered = "\n".join(o._repr_html_() for o in outputs if hasattr(o, "_repr_html_"))
    assert "LLM_DB_PASSWORD" in rendered and "not set" in rendered
    assert "pgvector is unreachable" in rendered

    tables = [t for out in outputs for t in _tables_in(out)]
    family_tables = [t for t in tables if "family:*" in t.text]
    assert family_tables, "expected the bib 'family:*' tag table to render"
    rows = {r["name"]: r["count"] for r in family_tables[0].data.to_dicts()}
    assert rows.get("family:CCM", 0) > 0


@_live
def test_section_7b_tolerates_missing_llm_deps(monkeypatch):
    """#720: the notebooks container installs the workspace without the
    ``llm`` extra. When ``llm.lineage`` cannot be imported, section 7b
    must render an explanatory note, never a traceback — the rest of the
    notebook still runs."""
    import sys

    monkeypatch.setitem(sys.modules, "llm.lineage", None)  # ImportError on import
    monkeypatch.delenv("LLM_DB_PASSWORD", raising=False)
    mod = _load()
    outputs, _defs = mod.app.run()
    rendered = "\n".join(o._repr_html_() for o in outputs if hasattr(o, "_repr_html_"))
    assert "7b. What the chat sees" in rendered
    assert "lineage module could not be imported here" in rendered
    assert "Traceback" not in rendered