tests/test_notebook.py

"""Lightweight structural tests for ``notebooks/train_kinchat.ipynb``.

We cannot execute the training notebook in CI — it requires a GPU, network
access to the deployed env, and ~45 min of wall clock. Instead we assert
that the notebook is a *valid* nbformat v4 file, that it contains the
marker strings the hackathon judges look for (GRPO, Unsloth, Qwen2.5, the
env URL, the four rubric names, the three plot filenames), that it has
enough cells to hang the required flow off, and that every code cell
parses as Python.

Keep these cheap; we just need a tripwire if a future edit breaks
structure.
"""
from __future__ import annotations

import ast
from pathlib import Path

import nbformat
import pytest


NOTEBOOK_PATH = (
    Path(__file__).resolve().parents[1] / "notebooks" / "train_kinchat.ipynb"
)


@pytest.fixture(scope="module")
def notebook():
    assert NOTEBOOK_PATH.exists(), f"missing notebook: {NOTEBOOK_PATH}"
    nb = nbformat.read(NOTEBOOK_PATH, as_version=4)
    return nb


def test_notebook_is_valid_nbformat_v4(notebook):
    assert notebook.nbformat == 4
    # nbformat.validate raises ValidationError on anything malformed.
    nbformat.validate(notebook)


def test_notebook_has_at_least_twelve_cells(notebook):
    assert len(notebook.cells) >= 12, (
        f"expected >=12 cells, got {len(notebook.cells)}"
    )


def test_notebook_has_both_markdown_and_code_cells(notebook):
    code_cells = [c for c in notebook.cells if c.cell_type == "code"]
    md_cells = [c for c in notebook.cells if c.cell_type == "markdown"]
    assert code_cells, "no code cells in notebook"
    assert md_cells, "no markdown cells in notebook"


def _all_source(notebook) -> str:
    """Concatenate every cell source into one blob for marker checks."""
    return "\n\n".join(c.source for c in notebook.cells)


@pytest.mark.parametrize(
    "marker",
    [
        "GRPO",
        "unsloth",
        "Qwen2.5",
        "reward_curve",
        "per_rubric_curves",
        "session_trust",
        "vex-0-kinchat.hf.space",
        # four rubric names:
        "leak",
        "audience_fit",
        "restraint",
        "trust_delta",
    ],
)
def test_notebook_contains_marker(notebook, marker):
    blob = _all_source(notebook)
    assert marker in blob, f"marker {marker!r} missing from notebook source"


def test_notebook_references_docs_plots_dir(notebook):
    blob = _all_source(notebook)
    # The plot cell must write into docs/plots/ relative to repo root.
    assert "docs/plots" in blob, "plot cell should target docs/plots/"


def test_notebook_mentions_huggingface_hub_push(notebook):
    blob = _all_source(notebook)
    # A Hub upload is required by the hackathon checklist.
    assert "huggingface_hub" in blob or "HfApi" in blob, (
        "notebook should push the LoRA to the HF Hub"
    )


def test_every_code_cell_parses_as_python(notebook):
    """Every code cell must be syntactically valid Python, ignoring
    notebook-specific magics (``%pip``, ``!cmd``, ``?ident``).
    """
    for i, cell in enumerate(notebook.cells):
        if cell.cell_type != "code":
            continue
        src = cell.source
        lines = src.splitlines()
        # Strip magics/shell/inspect lines; keep a raw blank so line numbers
        # in any future SyntaxError match the original cell.
        cleaned = []
        for line in lines:
            stripped = line.lstrip()
            if (
                stripped.startswith("%")
                or stripped.startswith("!")
                or stripped.startswith("?")
            ):
                cleaned.append("")
            else:
                cleaned.append(line)
        cleaned_src = "\n".join(cleaned)
        if not cleaned_src.strip():
            continue
        try:
            ast.parse(cleaned_src)
        except SyntaxError as exc:  # pragma: no cover — test failure path
            pytest.fail(f"cell {i} failed to parse: {exc}\n---\n{src}\n---")


def test_notebook_uses_grpo_trainer_or_documents_fallback(notebook):
    blob = _all_source(notebook)
    assert "GRPOTrainer" in blob, "notebook must reference GRPOTrainer (TRL)"
    # The fallback path must also be documented so graders see both rails.
    assert "manual GRPO" in blob.lower() or "fallback" in blob.lower(), (
        "notebook should document a fallback path if GRPOTrainer fails"
    )