tests/test_pipeline_structure.py

"""Structural tests for the training pipeline.

These do NOT invoke stages that need a GPU (SFT, gate inference, GRPO, eval
inference). They verify:

    * All stage modules are importable.
    * The stage entry-point functions exist with the expected names.
    * ``build_gate_prompts`` from stage 2 produces the right number of
      varied prompts (CPU-only).
    * The pipeline orchestrator's CLI parser accepts the documented flags.
    * The scripted eval policy in stage 4 works against the env (CPU-only).
"""
from __future__ import annotations

import importlib
from pathlib import Path


STAGE_MODULES = [
    "training.stages.stage_1_sft",
    "training.stages.stage_2_gate",
    "training.stages.stage_3_grpo",
    "training.stages.stage_4_eval",
]


def test_all_stage_modules_importable():
    """If any import fails (typo, missing dep, circular import), the whole
    pipeline is broken. Catch it here before we burn GPU."""
    for mod_name in STAGE_MODULES:
        # Stages depend on unsloth; we can still import-check if unsloth is
        # installed locally. If it's not, skip cleanly — the HF Space has it.
        try:
            importlib.import_module(mod_name)
        except ImportError as exc:
            if "unsloth" in str(exc).lower():
                import pytest
                pytest.skip(f"unsloth not available locally: {exc}")
            raise


def test_stage_entry_points_exist():
    """Each stage must expose a callable ``run_<stage>`` so pipeline.py
    can invoke it programmatically."""
    try:
        import training.stages.stage_1_sft as s1
        import training.stages.stage_2_gate as s2
        import training.stages.stage_3_grpo as s3
        import training.stages.stage_4_eval as s4
    except ImportError as exc:
        if "unsloth" in str(exc).lower():
            import pytest
            pytest.skip("unsloth not available locally")
        raise

    assert callable(s1.run_sft)
    assert callable(s2.run_gate)
    assert callable(s3.run_grpo)
    assert callable(s4.run_eval)


def test_gate_prompts_build_deterministically():
    """Gate prompts should be deterministic and diverse."""
    try:
        from training.stages.stage_2_gate import build_gate_prompts
    except ImportError as exc:
        if "unsloth" in str(exc).lower():
            import pytest
            pytest.skip("unsloth not available locally")
        raise
    a = build_gate_prompts()
    b = build_gate_prompts()
    assert len(a) == 20  # 4 tasks × 5 per task
    assert len(b) == 20
    # Deterministic across invocations
    assert [p["seed"] for p in a] == [p["seed"] for p in b]
    # All four tech tasks represented
    assert len({p["task_id"] for p in a}) == 4


def test_scripted_eval_policy_runs_on_env():
    """Stage 4's scripted baseline must produce valid parseable output."""
    try:
        from training.stages.stage_4_eval import _scripted_policy
    except ImportError as exc:
        if "unsloth" in str(exc).lower():
            import pytest
            pytest.skip("unsloth not available locally")
        raise
    from permanence.env import PermanenceEnv

    env = PermanenceEnv(config={"force_task": "task_log_cleanup"})
    obs, _ = env.reset(seed=100)
    completion = _scripted_policy(obs["text"])
    assert "<action" in completion
    assert "<reversibility" in completion


def test_pipeline_orchestrator_has_expected_stages():
    try:
        from training.pipeline import STAGES
    except ImportError as exc:
        if "unsloth" in str(exc).lower():
            import pytest
            pytest.skip("unsloth not available locally")
        raise
    assert STAGES == ["sft", "gate", "grpo", "eval"]


def test_reward_pack_usable_in_trl_shape():
    """TRL requires each reward func to accept (completions, **kwargs) and
    return list[float] of the same length."""
    from training.rewards import build_reward_pack
    pack = build_reward_pack(total_episodes=100)
    completions = [
        '<action id="fs_ls"/><reversibility level="R1" confidence="0.9"/>',
        "some bad output",
    ]
    for fn in pack.funcs:
        out = fn(completions, actual_r_levels=[1, 4], task_id=["task_x", "task_y"], seed=[1, 2])
        assert isinstance(out, list)
        assert len(out) == len(completions)
        assert all(isinstance(x, float) for x in out)