overlay/tests/test_hydra_modular.py

"""
Regression tests for W1's modularisation of train.py into the hydra/ package.

These tests verify that after modularisation:
  - The expected public symbols are importable from the stated sub-modules.
  - PostSemClawConfig instantiates with default args.
  - PostSemClawModel can be constructed, initialised, and produces a scalar
    loss on tiny inputs (batch=1, seq=32) without error.
  - train.py at the repo root is still importable as a Python module (i.e.
    the training-loop body is gated on ``if __name__ == "__main__":`` so a
    plain ``import`` doesn't execute it).
  - train.py is under 150 lines after modularisation (the main motiviation for
    W1's work is a thin orchestrator script, not a 900-line monolith).

If the hydra/ package does not exist yet (W1 is still running), every test in
this file is gracefully skipped so the test suite remains green.

Run:
    cd /home/mikeb/work/feather
    .venv/bin/pytest tests/test_hydra_modular.py -v
"""

import importlib
import os
import subprocess
import sys
import types
import pytest

# ---------------------------------------------------------------------------
# Module-level skip: hydra/ must exist as an importable package.
# pytest.importorskip cannot be used at module level without allow_module_level,
# and it doesn't work for relative paths.  We do the check manually.
# ---------------------------------------------------------------------------

_REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_HYDRA_INIT = os.path.join(_REPO, "hydra", "__init__.py")

if not os.path.isfile(_HYDRA_INIT):
    pytest.skip(
        "hydra/ package not found — W1 modularisation not yet complete. "
        "Re-run after hydra/__init__.py exists.",
        allow_module_level=True,
    )

# ---------------------------------------------------------------------------
# Helper: add repo root to sys.path so `import hydra` resolves to the local
# package, not the Apache Hydra framework if installed.
# ---------------------------------------------------------------------------

if _REPO not in sys.path:
    sys.path.insert(0, _REPO)


# ---------------------------------------------------------------------------
# Fixture: ensure 'prepare' stub is available so any transitive imports from
# train.py or hydra/ that do `from prepare import ...` don't crash.
# ---------------------------------------------------------------------------

def _ensure_prepare_stub():
    if "prepare" not in sys.modules:
        fake = types.ModuleType("prepare")
        fake.MAX_SEQ_LEN = 2048
        fake.TIME_BUDGET = 300
        fake.Tokenizer = object
        fake.make_dataloader = lambda *a, **kw: None
        fake.evaluate_bpb = lambda *a, **kw: 0.0
        sys.modules["prepare"] = fake


_ensure_prepare_stub()


# ---------------------------------------------------------------------------
# Test 1: public API is importable from the correct sub-modules
# ---------------------------------------------------------------------------

class TestHydraPublicAPI:
    def test_config_importable(self):
        """PostSemClawConfig is importable from hydra.config."""
        mod = importlib.import_module("hydra.config")
        assert hasattr(mod, "PostSemClawConfig"), (
            "hydra.config does not export PostSemClawConfig"
        )

    def test_model_importable(self):
        """PostSemClawModel is importable from hydra.model."""
        mod = importlib.import_module("hydra.model")
        assert hasattr(mod, "PostSemClawModel"), (
            "hydra.model does not export PostSemClawModel"
        )

    def test_optimizer_importable(self):
        """MuonAdamW is importable from hydra.optimizer."""
        mod = importlib.import_module("hydra.optimizer")
        assert hasattr(mod, "MuonAdamW"), (
            "hydra.optimizer does not export MuonAdamW"
        )

    def test_engram_importable(self):
        """GPUEngram is importable from hydra.engram (if Engram is top-level)."""
        try:
            mod = importlib.import_module("hydra.engram")
        except ImportError:
            pytest.skip("hydra.engram module does not exist — may be merged into hydra.model")
        assert hasattr(mod, "GPUEngram"), (
            "hydra.engram does not export GPUEngram"
        )


# ---------------------------------------------------------------------------
# Test 2: PostSemClawConfig default construction
# ---------------------------------------------------------------------------

class TestPostSemClawConfig:
    def test_default_instantiation(self):
        """PostSemClawConfig() should instantiate with all defaults."""
        from hydra.config import PostSemClawConfig  # noqa: PLC0415
        cfg = PostSemClawConfig()
        # Verify a few required fields exist and have sane defaults
        assert hasattr(cfg, "d_model"), "PostSemClawConfig missing d_model field"
        assert hasattr(cfg, "n_layer"), "PostSemClawConfig missing n_layer field"
        assert hasattr(cfg, "vocab_size"), "PostSemClawConfig missing vocab_size field"
        assert cfg.d_model > 0
        assert cfg.n_layer > 0
        assert cfg.vocab_size > 0

    def test_custom_instantiation(self):
        """PostSemClawConfig accepts keyword overrides."""
        from hydra.config import PostSemClawConfig  # noqa: PLC0415
        cfg = PostSemClawConfig(d_model=64, n_layer=2)
        assert cfg.d_model == 64
        assert cfg.n_layer == 2


# ---------------------------------------------------------------------------
# Test 3: PostSemClawModel forward pass with tiny inputs
# ---------------------------------------------------------------------------

class TestPostSemClawModelForward:
    @pytest.fixture
    def tiny_model(self):
        """Construct a tiny PostSemClawModel on CPU."""
        import torch  # noqa: PLC0415
        from hydra.config import PostSemClawConfig  # noqa: PLC0415
        from hydra.model import PostSemClawModel  # noqa: PLC0415

        # Use the smallest possible config that exercises all code paths.
        cfg = PostSemClawConfig(
            sequence_len=32,
            vocab_size=64,
            n_layer=2,
            d_model=32,
            d_state=8,
            headdim=16,
            n_heads=2,
            expand=2,
            engram_n_columns=16,
            engram_key_dim=8,
            engram_layer_idx=0,
            sdr_n_bits=128,
            sdr_target_active=3,
            sdr_delta_rank=4,
            htm_n_columns=32,
            htm_cells_per_column=4,
        )
        model = PostSemClawModel(cfg)
        model.init_weights()
        model.eval()
        return model

    def test_forward_returns_scalar_loss(self, tiny_model):
        """model(x, y, reduction='mean') returns a scalar loss."""
        import torch  # noqa: PLC0415

        B, T = 1, 32
        vocab = tiny_model.config.vocab_size
        idx = torch.randint(0, vocab, (B, T))
        targets = torch.randint(0, vocab, (B, T))

        with torch.no_grad():
            loss = tiny_model(idx, targets, reduction="mean")

        assert isinstance(loss, torch.Tensor), "forward did not return a tensor"
        assert loss.ndim == 0, f"expected scalar loss, got shape {loss.shape}"
        assert torch.isfinite(loss), f"loss is not finite: {loss.item()}"

    def test_forward_returns_per_token_loss(self, tiny_model):
        """model(x, y, reduction='none') returns (B*T,) per-token losses."""
        import torch  # noqa: PLC0415

        B, T = 1, 32
        vocab = tiny_model.config.vocab_size
        idx = torch.randint(0, vocab, (B, T))
        targets = torch.randint(0, vocab, (B, T))

        with torch.no_grad():
            losses = tiny_model(idx, targets, reduction="none")

        assert losses.shape == (B * T,), (
            f"expected shape ({B * T},), got {losses.shape}"
        )
        assert torch.all(torch.isfinite(losses)), "some per-token losses are not finite"


# ---------------------------------------------------------------------------
# Test 4: train.py at repo root is still importable (body gated on __main__)
# ---------------------------------------------------------------------------

class TestTrainPyImportable:
    def test_train_py_importable_as_module(self):
        """
        train.py must be importable without executing the training loop.
        We verify this by running `python -c "import importlib.util; ..."` in a
        subprocess to get a clean interpreter state, avoiding interference from
        the test process's already-patched sys.modules.
        """
        train_path = os.path.join(_REPO, "train.py")
        assert os.path.isfile(train_path), f"train.py not found at {train_path}"

        check_script = (
            "import importlib.util, sys; "
            "sys.path.insert(0, repr(_REPO)); "
            "spec = importlib.util.spec_from_file_location('train', repr(train_path)); "
            "assert spec is not None, 'spec is None'"
        ).replace("repr(_REPO)", repr(_REPO)).replace("repr(train_path)", repr(train_path))

        result = subprocess.run(
            [sys.executable, "-c", check_script],
            capture_output=True,
            text=True,
            timeout=10,
        )
        # A non-zero exit only means the assert failed, not a parse error —
        # either way we surface stderr for diagnosis.
        assert result.returncode == 0, (
            f"train.py spec creation failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
        )

    def test_train_py_under_150_lines(self):
        """
        After modularisation, train.py should be a thin orchestrator < 150 lines.
        This asserts the structural goal: all heavy logic lives in hydra/*.
        """
        train_path = os.path.join(_REPO, "train.py")
        with open(train_path) as fh:
            lines = fh.readlines()
        assert len(lines) < 150, (
            f"train.py has {len(lines)} lines — expected < 150 after modularisation. "
            "Move model/optimizer/config definitions to hydra/ sub-modules."
        )