stack/tests/test_ast_coverage.py

"""AST-based coverage discovery for the stack source tree.

This module is the authoritative "what exists vs what is implemented"
oracle for the entire stack.  It walks every ``.py`` file under
``src/``, parses the AST, and classifies every callable into one of:

- **implemented** — body contains real logic
- **stub** — body raises ``NotImplementedError`` (or is ``...`` / ``pass``)
- **narwhalify** — decorated with ``@nw.narwhalify`` or ``@narwhalify``
- **expr** — an ``Expr(...)`` instantiation in a pipe module
- **pydantic** — a class that inherits from ``BaseModel``

Tests then assert structural invariants:

1. Every ``Expr`` must have ``name``, ``fn``, ``output``, and ``after``
   set (no bare defaults that silently break pipelines).
2. Every ``@nw.narwhalify`` function must have a docstring.
3. Every Pydantic model exposed in a public ``__init__`` must be
   importable without error.
4. The ratio of implemented-to-stub functions is tracked; a warning is
   emitted (not a hard failure) if it regresses below the baseline.
5. Every ``pipe/*.py`` module that defines ``pipeline`` must export a
   ``Pipeline`` instance (not ``None`` or a bare list).

A separate parametrized test is generated for every discovered
``@nw.narwhalify`` function so that the pytest output clearly
identifies which functions are stubs — making CI output actionable
rather than a single blob of text.
"""

from __future__ import annotations

import ast
import importlib
import inspect
import sys
from pathlib import Path
from typing import NamedTuple

import pytest

# ── project root ──────────────────────────────────────────────────────────────

SRC_ROOT = Path(__file__).parent.parent / "src"
assert SRC_ROOT.exists(), f"src/ not found at {SRC_ROOT}"


# ── AST helpers ───────────────────────────────────────────────────────────────


def _iter_python_files(root: Path) -> list[Path]:
    """Return all .py files under *root*, sorted for determinism."""
    return sorted(root.rglob("*.py"))


def _decorator_names(node: ast.FunctionDef | ast.AsyncFunctionDef) -> list[str]:
    """Return flattened decorator name strings for a function node."""
    names: list[str] = []
    for dec in node.decorator_list:
        if isinstance(dec, ast.Name):
            names.append(dec.id)
        elif isinstance(dec, ast.Attribute):
            # e.g.  nw.narwhalify  →  "nw.narwhalify"
            names.append(f"{ast.unparse(dec)}")
        elif isinstance(dec, ast.Call):
            func = dec.func
            if isinstance(func, ast.Name):
                names.append(func.id)
            elif isinstance(func, ast.Attribute):
                names.append(ast.unparse(func))
    return names


def _is_narwhalify(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
    decs = _decorator_names(node)
    return any(
        d in ("narwhalify", "nw.narwhalify") or d.endswith(".narwhalify") for d in decs
    )


def _body_is_stub(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
    """Return True if the function body is a stub (NotImplementedError / pass / ...)."""
    body = node.body

    # Strip leading docstring
    stmts = body
    if (
        stmts
        and isinstance(stmts[0], ast.Expr)
        and isinstance(stmts[0].value, ast.Constant)
    ):
        stmts = stmts[1:]

    if not stmts:
        return True  # empty body after docstring

    # Single statement
    if len(stmts) == 1:
        stmt = stmts[0]
        # `pass`
        if isinstance(stmt, ast.Pass):
            return True
        # `...`
        if (
            isinstance(stmt, ast.Expr)
            and isinstance(stmt.value, ast.Constant)
            and stmt.value.value is ...
        ):
            return True
        # `raise NotImplementedError` or `raise NotImplementedError(...)`
        if isinstance(stmt, ast.Raise):
            exc = stmt.exc
            if exc is None:
                return False
            if isinstance(exc, ast.Name) and exc.id == "NotImplementedError":
                return True
            if isinstance(exc, ast.Call):
                func = exc.func
                if isinstance(func, ast.Name) and func.id == "NotImplementedError":
                    return True
                if (
                    isinstance(func, ast.Attribute)
                    and func.attr == "NotImplementedError"
                ):
                    return True
    return False


def _has_docstring(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
    body = node.body
    return (
        bool(body)
        and isinstance(body[0], ast.Expr)
        and isinstance(body[0].value, ast.Constant)
        and isinstance(body[0].value.value, str)
    )


def _base_names(node: ast.ClassDef) -> list[str]:
    """Return flattened base class name strings for a class node."""
    names: list[str] = []
    for base in node.bases:
        if isinstance(base, ast.Name):
            names.append(base.id)
        elif isinstance(base, ast.Attribute):
            names.append(ast.unparse(base))
    return names


def _is_pydantic_model(node: ast.ClassDef) -> bool:
    bases = _base_names(node)
    return any(b in ("BaseModel", "SQLTable", "SQLModel") for b in bases)


# ── data structures ───────────────────────────────────────────────────────────


class FunctionRecord(NamedTuple):
    path: Path
    module: str  # dotted module name relative to src
    qualname: str  # class.method or plain name
    lineno: int
    is_narwhalify: bool
    is_stub: bool
    has_docstring: bool


class ClassRecord(NamedTuple):
    path: Path
    module: str
    name: str
    lineno: int
    is_pydantic: bool


class ExprRecord(NamedTuple):
    path: Path
    module: str
    lineno: int
    has_name: bool
    has_fn: bool
    has_output: bool
    has_after: bool


# ── discovery ─────────────────────────────────────────────────────────────────


def _path_to_module(path: Path) -> str:
    rel = path.relative_to(SRC_ROOT)
    parts = list(rel.with_suffix("").parts)
    if parts[-1] == "__init__":
        parts = parts[:-1]
    return ".".join(parts)


def _collect_functions(tree: ast.Module, path: Path) -> list[FunctionRecord]:
    """Walk the AST and collect all function/method definitions."""
    records: list[FunctionRecord] = []
    module = _path_to_module(path)

    class Visitor(ast.NodeVisitor):
        def __init__(self) -> None:
            self._class_stack: list[str] = []

        def visit_ClassDef(self, node: ast.ClassDef) -> None:
            self._class_stack.append(node.name)
            self.generic_visit(node)
            self._class_stack.pop()

        def _visit_func(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None:
            qualname = (
                ".".join(self._class_stack + [node.name])
                if self._class_stack
                else node.name
            )
            records.append(
                FunctionRecord(
                    path=path,
                    module=module,
                    qualname=qualname,
                    lineno=node.lineno,
                    is_narwhalify=_is_narwhalify(node),
                    is_stub=_body_is_stub(node),
                    has_docstring=_has_docstring(node),
                )
            )
            # recurse into nested functions / methods
            old = self._class_stack[:]
            self.generic_visit(node)
            self._class_stack[:] = old

        def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
            self._visit_func(node)

        def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
            self._visit_func(node)

    Visitor().visit(tree)
    return records


def _collect_classes(tree: ast.Module, path: Path) -> list[ClassRecord]:
    records: list[ClassRecord] = []
    module = _path_to_module(path)
    for node in ast.walk(tree):
        if isinstance(node, ast.ClassDef):
            records.append(
                ClassRecord(
                    path=path,
                    module=module,
                    name=node.name,
                    lineno=node.lineno,
                    is_pydantic=_is_pydantic_model(node),
                )
            )
    return records


def _expr_keyword_names(call: ast.Call) -> set[str]:
    return {kw.arg for kw in call.keywords if kw.arg is not None}


def _collect_exprs(tree: ast.Module, path: Path) -> list[ExprRecord]:
    """Find all Expr(...) call sites in pipe modules."""
    if "pipe" not in path.parts and path.name not in ("pipe.py",):
        return []

    records: list[ExprRecord] = []
    module = _path_to_module(path)

    for node in ast.walk(tree):
        if not isinstance(node, ast.Call):
            continue
        func = node.func
        is_expr_call = (isinstance(func, ast.Name) and func.id == "Expr") or (
            isinstance(func, ast.Attribute) and func.attr == "Expr"
        )
        if not is_expr_call:
            continue
        kws = _expr_keyword_names(node)
        records.append(
            ExprRecord(
                path=path,
                module=module,
                lineno=node.lineno,
                has_name="name" in kws,
                has_fn="fn" in kws,
                has_output="output" in kws,
                has_after="after" in kws,
            )
        )
    return records


# ── build the full index ──────────────────────────────────────────────────────


def _build_index() -> tuple[
    list[FunctionRecord],
    list[ClassRecord],
    list[ExprRecord],
]:
    all_funcs: list[FunctionRecord] = []
    all_classes: list[ClassRecord] = []
    all_exprs: list[ExprRecord] = []

    for py_file in _iter_python_files(SRC_ROOT):
        # skip cache dirs and egg-info
        if any(
            part.startswith("__pycache__") or part.endswith(".egg-info")
            for part in py_file.parts
        ):
            continue
        try:
            source = py_file.read_text(encoding="utf-8")
            tree = ast.parse(source, filename=str(py_file))
        except SyntaxError:
            continue

        all_funcs.extend(_collect_functions(tree, py_file))
        all_classes.extend(_collect_classes(tree, py_file))
        all_exprs.extend(_collect_exprs(tree, py_file))

    return all_funcs, all_classes, all_exprs


# Run once at collection time — fast (pure AST, no imports)
_ALL_FUNCS, _ALL_CLASSES, _ALL_EXPRS = _build_index()

_NARWHALIFY_FUNCS = [f for f in _ALL_FUNCS if f.is_narwhalify]
_STUB_FUNCS = [f for f in _ALL_FUNCS if f.is_stub]
_PYDANTIC_CLASSES = [c for c in _ALL_CLASSES if c.is_pydantic]


# ── parametrize helpers ───────────────────────────────────────────────────────


def _func_id(r: FunctionRecord) -> str:
    return f"{r.module}::{r.qualname}"


def _expr_id(r: ExprRecord) -> str:
    return f"{r.module}::L{r.lineno}"


# ── 1. narwhalify function inventory ─────────────────────────────────────────


@pytest.mark.parametrize(
    "record", _NARWHALIFY_FUNCS, ids=[_func_id(r) for r in _NARWHALIFY_FUNCS]
)
def test_narwhalify_has_docstring(record: FunctionRecord) -> None:
    """Every @nw.narwhalify function must document its contract.

    Docstrings are the primary source for the ``Expr.doc`` property and
    for pipeline graph introspection.  A missing docstring is a gap in
    the data dictionary.
    """
    assert record.has_docstring, (
        f"{record.module}::{record.qualname} (line {record.lineno}) "
        f"is decorated with @nw.narwhalify but has no docstring.\n"
        f"  file: {record.path}"
    )


@pytest.mark.stub
@pytest.mark.parametrize(
    "record", _NARWHALIFY_FUNCS, ids=[_func_id(r) for r in _NARWHALIFY_FUNCS]
)
def test_narwhalify_implementation_status(record: FunctionRecord) -> None:
    """Report which @nw.narwhalify functions are stubs vs implemented.

    Marked ``stub`` so the full test run can be filtered:
        pytest -m 'not stub'   # skip stub inventory
        pytest -m stub         # only show stub status
    """
    status = "STUB" if record.is_stub else "IMPLEMENTED"
    # This test always passes — it's an inventory report, not a gate.
    # To gate on implementation completeness use
    # test_implementation_ratio_above_baseline below.
    _ = status  # surfaced in pytest -v output via the parametrize id


# ── 2. implementation ratio ───────────────────────────────────────────────────

# Baseline: fraction of non-dunder, non-private functions that are implemented.
# Adjust upward as stubs are filled in — never lower it.
_IMPL_RATIO_BASELINE = 0.30


def test_implementation_ratio_above_baseline() -> None:
    """Implemented function ratio must not regress below the baseline.

    This is a ratchet: once you implement more functions you raise the
    baseline so the test suite stays honest about coverage trajectory.
    """
    public = [
        f
        for f in _ALL_FUNCS
        if not f.qualname.startswith("_") and not f.qualname.startswith("test_")
    ]
    if not public:
        pytest.skip("No public functions found")

    implemented = [f for f in public if not f.is_stub]
    ratio = len(implemented) / len(public)

    # Informational summary in the failure message
    stub_list = "\n".join(
        f"  {f.module}::{f.qualname} (L{f.lineno})" for f in public if f.is_stub
    )

    assert ratio >= _IMPL_RATIO_BASELINE, (
        f"Implementation ratio {ratio:.1%} is below baseline {_IMPL_RATIO_BASELINE:.1%}.\n"
        f"Stubs remaining ({len(public) - len(implemented)}/{len(public)}):\n{stub_list}"
    )


def test_stub_inventory_is_complete() -> None:
    """All stubs must live in src/ — none in tests/ or dev/."""
    test_stubs = [
        f for f in _STUB_FUNCS if "tests" in str(f.path) or "dev" in str(f.path)
    ]
    assert not test_stubs, (
        "Stub functions found outside src/ — these should be real implementations:\n"
        + "\n".join(f"  {f.path}:{f.lineno} {f.qualname}" for f in test_stubs)
    )


# ── 3. Expr structural invariants ─────────────────────────────────────────────


@pytest.mark.parametrize("record", _ALL_EXPRS, ids=[_expr_id(r) for r in _ALL_EXPRS])
def test_expr_has_name(record: ExprRecord) -> None:
    """Every Expr(...) call must supply a ``name`` keyword argument."""
    assert record.has_name, (
        f"Expr at {record.module} line {record.lineno} is missing ``name=``.\n"
        f"  file: {record.path}"
    )


@pytest.mark.parametrize("record", _ALL_EXPRS, ids=[_expr_id(r) for r in _ALL_EXPRS])
def test_expr_has_fn(record: ExprRecord) -> None:
    """Every Expr(...) call must supply a ``fn`` keyword argument."""
    assert record.has_fn, (
        f"Expr at {record.module} line {record.lineno} is missing ``fn=``.\n"
        f"  file: {record.path}"
    )


@pytest.mark.parametrize("record", _ALL_EXPRS, ids=[_expr_id(r) for r in _ALL_EXPRS])
def test_expr_has_output(record: ExprRecord) -> None:
    """Every Expr(...) call must declare its ``output`` SQLTable contract."""
    assert record.has_output, (
        f"Expr at {record.module} line {record.lineno} is missing ``output=``.\n"
        f"Without an output contract, pipeline schema validation is impossible.\n"
        f"  file: {record.path}"
    )


@pytest.mark.parametrize("record", _ALL_EXPRS, ids=[_expr_id(r) for r in _ALL_EXPRS])
def test_expr_has_after(record: ExprRecord) -> None:
    """Every Expr(...) call must declare its ``after`` dependency list.

    An empty list (``after=[]``) is valid — the runner derives ordering
    from the function signature.  The invariant is that the author made
    an explicit choice, not that they left it to the default.
    """
    assert record.has_after, (
        f"Expr at {record.module} line {record.lineno} is missing ``after=``.\n"
        f"Declare ``after=[]`` explicitly if there are no upstream dependencies.\n"
        f"  file: {record.path}"
    )


# ── 4. Pydantic model importability ──────────────────────────────────────────


def _public_pydantic_models() -> list[ClassRecord]:
    """Pydantic models in public modules (not _private or test files)."""
    return [
        c
        for c in _PYDANTIC_CLASSES
        if not any(part.startswith("_") for part in c.module.split("."))
        and "test" not in c.module
        and "dev" not in c.module
    ]


_PUBLIC_MODELS = _public_pydantic_models()


@pytest.mark.parametrize(
    "record",
    _PUBLIC_MODELS,
    ids=[f"{c.module}::{c.name}" for c in _PUBLIC_MODELS],
)
def test_pydantic_model_importable(record: ClassRecord) -> None:
    """Every public Pydantic model must be importable without error.

    Import failures at test time surface missing dependencies, circular
    imports, or broken ``__init__.py`` re-exports early.
    """
    # Insert src onto path if needed
    src_str = str(SRC_ROOT)
    if src_str not in sys.path:
        sys.path.insert(0, src_str)

    try:
        mod = importlib.import_module(record.module)
    except ImportError as exc:
        pytest.fail(
            f"Cannot import {record.module} (needed for {record.name}): {exc}\n"
            f"  file: {record.path}"
        )

    cls = getattr(mod, record.name, None)
    if cls is None:
        # Class may be private or re-exported under a different name — skip
        pytest.skip(f"{record.name} not directly accessible on {record.module}")

    if not isinstance(cls, type):
        pytest.skip(f"{record.name} is not a class on {record.module}")

    # Basic sanity: must have model_fields (Pydantic v2 class attribute)
    assert hasattr(cls, "model_fields"), (
        f"{record.module}.{record.name} does not look like a Pydantic v2 model "
        f"(missing ``model_fields`` class attribute)."
    )


# ── 5. Pipeline module structural checks ──────────────────────────────────────


def _find_pipeline_modules() -> list[Path]:
    """Find all pipe/*.py files that assign a ``pipeline`` name."""
    results: list[Path] = []
    for path in _iter_python_files(SRC_ROOT):
        if "pipe" not in path.parts:
            continue
        if path.name.startswith("_"):
            continue
        try:
            source = path.read_text(encoding="utf-8")
            tree = ast.parse(source)
        except (SyntaxError, OSError):
            continue
        for node in ast.walk(tree):
            if isinstance(node, (ast.Assign, ast.AnnAssign)) and _assigns_name(
                node, "pipeline"
            ):
                results.append(path)
                break
    return results


def _assigns_name(node: ast.Assign | ast.AnnAssign, name: str) -> bool:
    if isinstance(node, ast.Assign):
        return any((isinstance(t, ast.Name) and t.id == name) for t in node.targets)
    if isinstance(node, ast.AnnAssign):
        return isinstance(node.target, ast.Name) and node.target.id == name
    return False


_PIPELINE_MODULES = _find_pipeline_modules()


@pytest.mark.parametrize(
    "path",
    _PIPELINE_MODULES,
    ids=[_path_to_module(p) for p in _PIPELINE_MODULES],
)
def test_pipeline_module_exports_pipeline_instance(path: Path) -> None:
    """Each pipe module that defines ``pipeline`` must export a Pipeline.

    A ``None`` or list value would silently break the runner.
    """
    src_str = str(SRC_ROOT)
    if src_str not in sys.path:
        sys.path.insert(0, src_str)

    module_name = _path_to_module(path)
    try:
        mod = importlib.import_module(module_name)
    except ImportError as exc:
        pytest.fail(f"Cannot import {module_name}: {exc}")

    pipeline = getattr(mod, "pipeline", None)
    assert pipeline is not None, (
        f"{module_name} defines ``pipeline`` but the attribute is None at import time."
    )
    # Check it has an ``exprs`` attribute (duck-type Pipeline check)
    assert hasattr(pipeline, "exprs"), (
        f"{module_name}.pipeline does not look like a Pipeline — "
        f"missing ``.exprs`` attribute.  Got: {type(pipeline)}"
    )
    assert isinstance(pipeline.exprs, list), (
        f"{module_name}.pipeline.exprs must be a list, got {type(pipeline.exprs)}"
    )
    assert len(pipeline.exprs) > 0, (
        f"{module_name}.pipeline.exprs is empty — no expressions registered."
    )


# ── 6. param_to_table naming convention ──────────────────────────────────────


def _collect_narwhalify_param_names() -> list[tuple[str, str, str]]:
    """Return (module, qualname, param_name) for every narwhalify param."""
    results: list[tuple[str, str, str]] = []
    src_str = str(SRC_ROOT)
    if src_str not in sys.path:
        sys.path.insert(0, src_str)

    for rec in _NARWHALIFY_FUNCS:
        try:
            mod = importlib.import_module(rec.module)
            fn = getattr(mod, rec.qualname.split(".")[0], None)
            if fn is None or not callable(fn):
                continue
            sig = inspect.signature(fn)
            for pname in sig.parameters:
                results.append((rec.module, rec.qualname, pname))
        except Exception:
            continue
    return results


def test_narwhalify_param_names_follow_convention() -> None:
    """Parameter names of @nw.narwhalify functions must follow the __ convention.

    Valid forms:
    - ``df``                   — single unnamed input (passthrough)
    - ``schema__table``        → ``schema.table``
    - ``schema___private``     → ``schema._private``
    - bare single word like ``rvu``, ``gpci`` (for PFS-style positional args)

    Invalid: mixed case, triple underscores at start, or ``__`` prefix
    that would resolve to ``.table`` (missing schema).
    """
    violations: list[str] = []
    params = _collect_narwhalify_param_names()

    for module, qualname, pname in params:
        # Skip self/cls
        if pname in ("self", "cls"):
            continue
        # df or single-word bare name — always fine
        if "__" not in pname:
            continue
        # Must have a non-empty schema part before the first __
        parts = pname.split("__", 1)
        schema = parts[0]
        if not schema:
            violations.append(
                f"  {module}::{qualname} param ``{pname}`` starts with __ "
                f"(missing schema prefix)"
            )

    assert not violations, (
        "Parameter naming convention violations found:\n" + "\n".join(violations)
    )


# ── 7. no bare ``raise NotImplementedError`` in __init__.py ──────────────────


def test_no_stubs_in_init_files() -> None:
    """__init__.py files must not contain NotImplementedError stubs.

    Init files declare the public API surface.  A stub there breaks any
    consumer that imports from the package.
    """
    bad: list[str] = []
    for rec in _STUB_FUNCS:
        if rec.path.name == "__init__.py":
            bad.append(f"  {rec.path}:{rec.lineno}  {rec.qualname}")
    assert not bad, "Stub functions found in __init__.py files:\n" + "\n".join(bad)


# ── 8. express / pipe module symmetry ────────────────────────────────────────


def test_express_pipe_module_symmetry() -> None:
    """Each aco/express/<module>.py should have a matching aco/pipe/<module>.py.

    The express layer holds narwhalify functions; the pipe layer assembles them
    into Pipeline objects.  Asymmetry usually means a module was added to one
    side but not wired up on the other.
    """
    express_dir = SRC_ROOT / "aco" / "express"
    pipe_dir = SRC_ROOT / "aco" / "pipe"

    if not express_dir.exists() or not pipe_dir.exists():
        pytest.skip("aco/express or aco/pipe not found")

    express_mods = {
        p.stem for p in express_dir.glob("*.py") if not p.stem.startswith("_")
    }
    pipe_mods = {
        p.stem
        for p in pipe_dir.glob("*.py")
        if not p.stem.startswith("_") and p.stem not in ("runner", "opps")
    }

    missing_in_pipe = express_mods - pipe_mods
    missing_in_express = pipe_mods - express_mods

    msgs: list[str] = []
    if missing_in_pipe:
        msgs.append(
            "express/ modules with no matching pipe/ module:\n"
            + "\n".join(f"  aco/express/{m}.py" for m in sorted(missing_in_pipe))
        )
    if missing_in_express:
        msgs.append(
            "pipe/ modules with no matching express/ module:\n"
            + "\n".join(f"  aco/pipe/{m}.py" for m in sorted(missing_in_express))
        )

    assert not msgs, "\n".join(msgs)


# ── 9. Tag model factory method coverage ─────────────────────────────────────


def test_tag_factory_methods_return_tag_instances() -> None:
    """All Tag.* factory methods must return Tag instances with correct namespace."""
    src_str = str(SRC_ROOT)
    if src_str not in sys.path:
        sys.path.insert(0, src_str)

    from bib.tag import Tag

    cases = [
        (Tag.module("aco"), "module", "aco"),
        (Tag.table("pfs.rvu"), "table", "pfs.rvu"),
        (Tag.source("cms-website"), "source", "cms-website"),
        (Tag.year(2026), "year", "2026"),
        (Tag.rule("cms-1807-f"), "rule", "cms-1807-f"),
        (Tag.file("rvu"), "file", "rvu"),
        (Tag.sup("parent-key"), "sup", "parent-key"),
    ]

    for tag, expected_ns, expected_val in cases:
        assert isinstance(tag, Tag), f"Expected Tag, got {type(tag)}"
        assert tag.namespace == expected_ns, (
            f"Tag.{expected_ns}({expected_val!r}).namespace = {tag.namespace!r}, "
            f"expected {expected_ns!r}"
        )
        assert tag.value == expected_val, (
            f"Tag.{expected_ns}({expected_val!r}).value = {tag.value!r}, "
            f"expected {expected_val!r}"
        )