giant/tests/test_config_consumed_keys.py

"""Consumed-keys audit (issues.md Issue 5).

`validate_config_keys` (`giant/config.py`) only checks that a config key is
*declared* — present somewhere in `DEFAULT_CONFIG`, which is generated from
the frozen dataclasses. It says nothing about whether anything actually
*reads* the value once parsed. Issues 1, 2 and 4 are three keys that slipped
through exactly that gap: declared, round-tripped, silently ignored. This
module walks every leaf path in `DEFAULT_CONFIG` and asserts each is either
genuinely consumed by the model-building/training/rollout code, or explicitly
recorded in `_KNOWN_UNUSED` with a reason.

"Consumed" is approximated by static analysis rather than true call-graph
reachability: for each leaf path's field name, does it appear anywhere in a
fixed whitelist of source files as a real attribute access, a dict-key-shaped
string constant, or a function/constructor parameter name (the last of these
because `Router` subclasses receive their config via `**kwargs` filtered by
signature — see `giant.model.routers.build_router`)? Docstrings are excluded
from the string-constant scan so prose mentioning a dotted config path in
passing can't masquerade as a read of it. This whitelist-based approach is
deliberately narrower than "anywhere in `giant/`": scanning the whole package
produces false negatives from unrelated identifier collisions (e.g.
`giant/analysis/router_gating.py`'s `_top1_shares(..., order: list, ...)`
parameter would otherwise make `stage2_model.autoregressive.order` read as
"consumed").
"""

import ast
from pathlib import Path

from giant.config import DEFAULT_CONFIG

_REPO_ROOT = Path(__file__).resolve().parents[1]

# Files that legitimately consume model_config / training config at
# build/train/rollout time. Not `giant/cli.py` (a CLI flag existing is not
# consumption — that's precisely how Issue 1 slipped through), not
# `giant/config.py` itself (declaring/parsing a field is not reading it), and
# not `giant/model/_legacy.py` (the protected v0.2 migration surface, which
# intentionally re-derives old flat keys under old names).
_CONSUMER_ROOTS = ("giant/model", "giant/training")
_CONSUMER_FILES = (
    "giant/sample.py",
    "giant/pipeline.py",
    "giant/rollout.py",
    "giant/checkpoint_io.py",
    "giant/particles.py",
    "giant/materials.py",
)
_EXCLUDED_FILES = ("giant/model/_legacy.py",)

# Leaf DEFAULT_CONFIG paths that are declared but not (yet) read anywhere in
# the consumer whitelist above. Each entry must name the issue that tracks
# it. If a key here starts showing up as consumed, the fix landed and this
# entry is stale — see test_known_unused_allow_list_has_no_stale_entries.
_KNOWN_UNUSED = {
    "stage2_model.stage1_context": (
        "issues.md Issue 1 — trainers.py hardcodes stage1_ctx to the "
        "ground-truth stage-1 output; 'sampled' is now rejected loudly by "
        "validate_config (not silently accepted), but the key still isn't "
        "read by any build/train consumer file since only 'truth' can pass "
        "validation — see Issue 16 for the real implementation"
    ),
    "stage1_model.wgan.critic_hidden_dim": (
        "issues.md Issue 2 — build_critics always sizes the critic off the generator's own hidden_dim, never this key"
    ),
    "stage1_model.wgan.critic_n_res_blocks": ("issues.md Issue 2 — same as critic_hidden_dim"),
    "stage2_model.wgan.critic_hidden_dim": ("issues.md Issue 2 — same as stage1_model.wgan.critic_hidden_dim"),
    "stage2_model.wgan.critic_n_res_blocks": ("issues.md Issue 2 — same as stage1_model.wgan.critic_hidden_dim"),
    "stage2_model.autoregressive.order": (
        "issues.md Issue 4 — validate_config checks history/teacher_forcing "
        "but never order, and nothing reads it either"
    ),
}

# "lambda" is a Python keyword, so the dataclasses expose the dict key
# "lambda" as the field `lambda_weight` (giant/config.py:49-50).
_FIELD_NAME_OVERRIDES = {"lambda": "lambda_weight"}


def _leaf_paths(node: dict, prefix: str = "") -> list[str]:
    paths = []
    for key, value in node.items():
        if prefix == "" and key == "meta":
            continue
        path = f"{prefix}.{key}" if prefix else key
        if isinstance(value, dict):
            paths.extend(_leaf_paths(value, path))
        else:
            paths.append(path)
    return paths


def _field_name(leaf_path: str) -> str:
    name = leaf_path.rsplit(".", 1)[-1]
    return _FIELD_NAME_OVERRIDES.get(name, name)


def _is_docstring_expr(expr: ast.Expr) -> bool:
    return isinstance(expr.value, ast.Constant) and isinstance(expr.value.value, str)


def _collect_names(source: str, filename: str) -> set[str]:
    tree = ast.parse(source, filename=filename)
    docstring_ids = set()
    for node in ast.walk(tree):
        if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
            body = getattr(node, "body", [])
            if body and isinstance(body[0], ast.Expr) and _is_docstring_expr(body[0]):
                docstring_ids.add(id(body[0].value))

    names: set[str] = set()
    for node in ast.walk(tree):
        if isinstance(node, ast.Attribute):
            names.add(node.attr)
        elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
            names.add(node.value)
        elif isinstance(node, ast.arg):
            names.add(node.arg)
        elif isinstance(node, ast.keyword) and node.arg is not None:
            names.add(node.arg)
    return names


def _consumer_files() -> list[Path]:
    files: set[Path] = {_REPO_ROOT / f for f in _CONSUMER_FILES}
    for root in _CONSUMER_ROOTS:
        files |= set((_REPO_ROOT / root).rglob("*.py"))
    files -= {_REPO_ROOT / f for f in _EXCLUDED_FILES}
    return sorted(files)


def _consumed_names() -> set[str]:
    names: set[str] = set()
    for path in _consumer_files():
        names |= _collect_names(path.read_text(), str(path))
    return names


def test_every_config_key_is_consumed_or_allow_listed():
    consumed = _consumed_names()
    unconsumed = {p for p in _leaf_paths(DEFAULT_CONFIG) if _field_name(p) not in consumed}
    unexplained = unconsumed - _KNOWN_UNUSED.keys()
    assert not unexplained, (
        f"config key(s) {sorted(unexplained)} are declared in DEFAULT_CONFIG "
        "but not read anywhere in the build/train/rollout consumer files "
        f"({[str(f.relative_to(_REPO_ROOT)) for f in _consumer_files()]}) — "
        "either wire the key up, or add it to _KNOWN_UNUSED with a reason "
        "(see issues.md Issue 5)"
    )


def test_known_unused_allow_list_has_no_stale_entries():
    consumed = _consumed_names()
    all_paths = set(_leaf_paths(DEFAULT_CONFIG))
    stale = {p for p in _KNOWN_UNUSED if p not in all_paths or _field_name(p) in consumed}
    assert not stale, (
        f"_KNOWN_UNUSED entry/entries {sorted(stale)} no longer belong on the "
        "allow-list — either the key was removed from DEFAULT_CONFIG, or it "
        "is now consumed (the underlying issue was fixed). Remove the stale "
        "entry/entries."
    )