Blame for tests/test_continuous_adapters.py

"""Contracts for the pinned LangFlow and RDLM continuous teachers."""

from dataclasses import replace
import hashlib
import json
import os
from pathlib import Path
import sys
from types import SimpleNamespace

import pytest
import yaml

import dlb.adapters.capture as capture_module
from dlb.adapters.base import AdapterError
from dlb.adapters.langflow import LangFlowAdapter
from dlb.adapters.rdlm import RDLMAdapter
from dlb.command import main as command_main
from dlb.registry import load_registry
from dlb.runner import RunRequest, _resolve_checkpoint_provenance


ROOT = Path(__file__).parents[1]


def request(model: str, dataset: str, steps: int = 32, samples: int = 17) -> RunRequest:
    return RunRequest(
        run_id=f"{model}-{dataset}-steps-{steps}",
        model_id=model,
        dataset_id=dataset,
        step_count=steps,
        seed=42,
        sample_count=samples,
    )


def run_dir(root: Path, item: RunRequest) -> Path:
    return (
        root
        / "results"
        / "samples"
        / item.dataset_id
        / item.model_id
        / f"steps_{item.step_count}"
    )


def option(command: list[str], key: str) -> str:
    index = command.index(key)
    return command[index + 1]


def override(command: list[str], key: str) -> str:
    prefix = key + "="
    values = [argument[len(prefix) :] for argument in command if argument.startswith(prefix)]
    assert len(values) == 1, (key, command)
    return values[0]


def prepare_conversion_root(tmp_path: Path) -> Path:
    (tmp_path / "artifacts").mkdir()
    (tmp_path / "configs").mkdir()
    for relative in (
        Path("artifacts/data.yaml"),
        Path("artifacts/checkpoints.yaml"),
        Path("configs/experiments.yaml"),
    ):
        target = tmp_path / relative
        target.write_text((ROOT / relative).read_text(encoding="utf-8"), encoding="utf-8")
    return tmp_path


def prepare_command_root(tmp_path: Path) -> Path:
    root = prepare_conversion_root(tmp_path)
    for relative in (
        Path("adapters/sample_langflow.py"),
        Path("upstreams/rdlm/main.py"),
    ):
        target = root / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_text("# fixture\n", encoding="utf-8")
    return root


def canonical_conversion_request(root: Path, item: RunRequest) -> RunRequest:
    langflow_resource = (
        "langflow_lm1b_hf" if item.dataset_id == "lm1b" else "langflow_owt_hf"
    )
    resources = {
        "langflow": (
            langflow_resource,
            "continuous_langflow",
            [
                "config.json",
                "config.py",
                "model.py",
                "model.safetensors",
            ],
        ),
        "rdlm": (
            "rdlm_lm1b_drive",
            "continuous_rdlm",
            [
                "LM1B/checkpoint.pth",
                "LM1B/config.yaml",
                "LM1B/sde.pkl",
                "Text8/checkpoint.pth",
                "Text8/config.yaml",
                "Text8/sde.pkl",
            ],
        ),
    }
    resource, family, relative_files = resources[item.model_id]
    manifest = root / "artifacts" / "checkpoints.yaml"
    manifest_sha256 = hashlib.sha256(manifest.read_bytes()).hexdigest()
    files = [
        {
            "path": f"checkpoints/fixture/{resource}/{relative}",
            "size_bytes": index + 1,
            "sha256": f"{index + 1:064x}",
        }
        for index, relative in enumerate(relative_files)
    ]
    (root / "artifacts" / "checkpoint_lock.json").write_text(
        json.dumps(
            {
                "manifest_sha256": manifest_sha256,
                "resources": {resource: {"status": "downloaded", "files": files}},
            }
        ),
        encoding="utf-8",
    )
    provenance = _resolve_checkpoint_provenance(root, item, None)
    return replace(
        item,
        checkpoint_sha256=provenance.sha256,
        checkpoint_lock_id=provenance.lock_id,
        checkpoint_selection=provenance.selection,
        checkpoint_teacher_family=provenance.teacher_family,
    )


def write_capture(path: Path, count: int, length: int) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(
            {
                "schema": "dlb-upstream-token-capture-v1",
                "samples": [
                    {
                        "sample_id": index,
                        "text": f"sample {index}",
                        "token_ids": [index + 1] * length,
                    }
                    for index in range(count)
                ],
            }
        ),
        encoding="utf-8",
    )


def prepare_tokenizer_binding(
    root: Path, dataset: str
) -> tuple[Path, Path, Path]:
    """Create the minimal immutable tokenizer lock consumed by capture tests."""

    data_config = root / "artifacts" / "data.yaml"
    data_config.parent.mkdir(parents=True, exist_ok=True)
    data_config.write_text(
        (ROOT / "artifacts/data.yaml").read_text(encoding="utf-8"),
        encoding="utf-8",
    )
    document = yaml.safe_load(data_config.read_text(encoding="utf-8"))
    tokenizer_id = document["datasets"][dataset]["tokenizer"]
    revision = document["models"][tokenizer_id]
    snapshot = (
        root
        / "locked-tokenizers"
        / f"models--{tokenizer_id}"
        / "snapshots"
        / revision
    )
    snapshot.mkdir(parents=True)
    downloads = root / "data/manifests/downloads.json"
    downloads.parent.mkdir(parents=True)
    downloads.write_text(
        json.dumps(
            {
                "schema_version": 1,
                "models": {
                    tokenizer_id: {
                        "repo_id": tokenizer_id,
                        "revision": revision,
                        "snapshot_path": snapshot.relative_to(root).as_posix(),
                        "size_bytes": 1,
                    }
                },
            }
        ),
        encoding="utf-8",
    )
    return data_config, downloads, snapshot


def test_registry_exposes_langflow_and_rdlm_supported_continuous_cells() -> None:
    """Catch supported continuous cells being silently substituted or rejected."""

    registry = load_registry(ROOT / "configs" / "experiments.yaml")

    assert registry.models["langflow"].datasets["owt"].status == "supported"
    assert registry.models["langflow"].datasets["lm1b"].status == "supported"
    assert registry.models["rdlm"].datasets["lm1b"].status == "supported"
    assert registry.models["rdlm"].datasets["owt"].status == "unsupported"


@pytest.mark.parametrize(
    ("dataset", "steps", "tokenizer_name", "tokenizer_revision", "sequence_length"),
    [
        ("lm1b", 32, "bert-base-uncased", "86b5e0934494bd15c9632b12f734a8a67f723594", 128),
        ("owt", 1024, "gpt2", "607a30d783dfa663caf39e06633721c8d4cfcd7e", 1024),
    ],
)
def test_langflow_renders_the_pinned_inference_argparse_contract(
    tmp_path: Path,
    dataset: str,
    steps: int,
    tokenizer_name: str,
    tokenizer_revision: str,
    sequence_length: int,
) -> None:
    """Catch flags that are not accepted by the pinned inference argparse parser."""

    root = prepare_conversion_root(tmp_path)
    wrapper = root / "adapters" / "sample_langflow.py"
    wrapper.parent.mkdir(parents=True)
    wrapper.write_text("# fixture\n", encoding="utf-8")
    entrypoint = root / "upstreams" / "langflow" / "inference.py"
    entrypoint.parent.mkdir(parents=True)
    entrypoint.write_text("# fixture\n", encoding="utf-8")
    item = request("langflow", dataset, steps=steps, samples=17)
    command = LangFlowAdapter().render_command(item, run_dir(root, item), dry_run=True)

    assert command[1:3] == ["-B", "-u"]
    assert Path(override(command, "--upstream-entrypoint")) == wrapper
    assert override(command, "--capture-kind") == "langflow"
    assert Path(override(command, "--data-config-path")) == root / "artifacts/data.yaml"
    assert Path(override(command, "--downloads-manifest-path")) == (
        root / "data/manifests/downloads.json"
    )
    assert override(command, "--dataset-id") == dataset
    assert Path(override(command, "--tokenizer-snapshot")) == (
        root
        / "data/raw/huggingface/hub"
        / f"models--{tokenizer_name}"
        / "snapshots"
        / tokenizer_revision
    )
    checkpoint = Path(option(command, "--checkpoint"))
    assert checkpoint == root / f"checkpoints/official/langflow/{dataset}/model.safetensors"
    assert option(command, "--num_samples") == "17"
    assert option(command, "--batch_size") == "1"
    assert option(command, "--num_steps") == str(steps)
    assert option(command, "--seq_length") == str(sequence_length)
    assert option(command, "--seed") == "42"
    assert Path(option(command, "--output")) == run_dir(root, item) / "upstream_samples.txt"
    assert option(command, "--tokenizer") == tokenizer_name
    assert not any(
        argument
        in {
            "--num-samples",
            "--num-steps",
            "--seq-length",
        }
        for argument in command
    )


def test_langflow_accepts_many_step_grid_values(
    tmp_path: Path,
) -> None:
    """Catch regressing LangFlow back to the fixed 1024-step-only cell."""

    root = prepare_conversion_root(tmp_path)
    wrapper = root / "adapters" / "sample_langflow.py"
    wrapper.parent.mkdir(parents=True)
    wrapper.write_text("# fixture\n", encoding="utf-8")
    item = request("langflow", "owt", steps=32, samples=17)

    command = LangFlowAdapter().render_command(item, run_dir(root, item), dry_run=True)

    assert option(command, "--num_steps") == "32"


def test_rdlm_renders_the_official_sde_sampler_and_saved_asset_trio(
    tmp_path: Path,
) -> None:
    """Catch RDLM bypassing its official Hydra sampler or omitting saved release assets."""

    root = prepare_command_root(tmp_path)
    item = request("rdlm", "lm1b", steps=1000, samples=17)
    command = RDLMAdapter().render_command(item, run_dir(root, item), dry_run=True)
    asset_root = root / "checkpoints/official/rdlm/lm1b/LM1B"

    assert command[1:3] == ["-B", "-u"]
    assert Path(override(command, "--upstream-entrypoint")) == root / "upstreams/rdlm/main.py"
    assert override(command, "--capture-kind") == "rdlm"
    assert Path(override(command, "--data-config-path")) == root / "artifacts/data.yaml"
    assert Path(override(command, "--downloads-manifest-path")) == (
        root / "data/manifests/downloads.json"
    )
    assert override(command, "--dataset-id") == "lm1b"
    assert Path(override(command, "--tokenizer-snapshot")) == (
        root
        / "data/raw/huggingface/hub/models--bert-base-uncased/snapshots"
        / "86b5e0934494bd15c9632b12f734a8a67f723594"
    )
    assert Path(override(command, "--saved-config-path")) == asset_root / "config.yaml"
    assert Path(override(command, "--saved-sde-path")) == asset_root / "sde.pkl"
    assert override(command, "--expected-samples") == "17"
    assert override(command, "run_mode") == "sample"
    assert override(command, "server") == "sample"
    assert override(command, "exp") == "sample_lm1b"
    assert Path(override(command, "model_path")) == asset_root / "checkpoint.pth"
    assert override(command, "sampling.predictor") == "grw"
    assert override(command, "sampling.steps") == "1000"
    assert override(command, "sampling.batch_per_gpu") == "8"
    assert override(command, "seed") == "42"
    assert override(command, "ngpus") == "1"
    assert override(command, "eval.entropy") == "False"
    assert override(command, "eval.nll") == "False"
    assert override(command, "eval.gen_ppl") == "False"


@pytest.mark.parametrize(
    ("adapter", "model", "dataset"),
    [
        (LangFlowAdapter(), "langflow", "owt"),
        (RDLMAdapter(), "rdlm", "lm1b"),
    ],
)
def test_real_render_requires_canonical_checkpoint_bytes(adapter, model: str, dataset: str) -> None:
    """Catch real mode accepting dry-run-only expected paths."""

    item = request(model, dataset, steps=1000 if model == "rdlm" else 32)
    with pytest.raises(AdapterError, match="checkpoint"):
        adapter.build_command(item, run_dir(ROOT, item))


@pytest.mark.parametrize(
    ("adapter", "model", "dataset", "length", "family"),
    [
        (LangFlowAdapter(), "langflow", "owt", 1024, "continuous_langflow"),
        (RDLMAdapter(), "rdlm", "lm1b", 128, "continuous_rdlm"),
    ],
)
def test_capture_conversion_preserves_exact_upstream_ids_and_latency_sentinel(
    tmp_path: Path, adapter, model: str, dataset: str, length: int, family: str
) -> None:
    """Catch lossy parsing, wrong lengths/counts, or 0.0 being presented as measured latency."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request(model, dataset, samples=17))
    output_dir = run_dir(root, item)
    write_capture(output_dir / "upstream_token_ids.json", 17, length)

    records = list(adapter.convert_outputs(item, output_dir))
    metadata = json.loads((output_dir / "conversion_metadata.json").read_text())

    assert len(records) == 17
    assert records[0].token_ids == [1] * length
    assert metadata["format"] == "dlb-upstream-token-capture-v1"
    assert metadata["token_ids_source"] == "upstream"
    assert metadata["generated_samples"] == 17
    assert metadata["trimmed_samples"] == 0
    assert metadata["teacher_family"] == family
    assert metadata["generation_seconds_source"] == "unavailable_excluded_sentinel"
    assert metadata["exclude_from_latency"] is True
    assert all(record.generation_seconds == 0.0 for record in records)


@pytest.mark.parametrize("changed_name", ["LM1B/config.yaml", "LM1B/sde.pkl"])
def test_rdlm_conversion_binds_saved_config_and_sde_checksums_before_output_reads(
    tmp_path: Path, changed_name: str
) -> None:
    """Catch provenance that hashes only checkpoint.pth and ignores the saved trio."""

    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("rdlm", "lm1b", samples=1))
    lock_path = root / "artifacts/checkpoint_lock.json"
    lock = json.loads(lock_path.read_text())
    files = lock["resources"]["rdlm_lm1b_drive"]["files"]
    target = next(file for file in files if file["path"].endswith("/" + changed_name))
    target["sha256"] = "f" * 64
    lock_path.write_text(json.dumps(lock), encoding="utf-8")
    output_dir = run_dir(root, item)
    output_dir.mkdir(parents=True)

    with pytest.raises(AdapterError, match="checkpoint SHA"):
        list(RDLMAdapter().convert_outputs(item, output_dir))


def test_rdlm_retokenizes_variable_upstream_rows_offline_and_records_why(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Catch variable post-BOS RDLM rows being padded without explicit tokenizer provenance."""

    class FakeTokenizer:
        eos_token = "[SEP]"
        pad_token = "[PAD]"
        pad_token_id = 0

        def __call__(self, texts, **settings):
            assert settings["max_length"] == 128
            return {"input_ids": [[11] + [0] * 127 for _ in texts]}

    class FakeAutoTokenizer:
        @staticmethod
        def from_pretrained(name, **settings):
            assert name == "bert-base-uncased"
            assert settings == {
                "revision": "86b5e0934494bd15c9632b12f734a8a67f723594",
                "local_files_only": True,
            }
            return FakeTokenizer()

    monkeypatch.setitem(
        sys.modules, "transformers", SimpleNamespace(AutoTokenizer=FakeAutoTokenizer)
    )
    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("rdlm", "lm1b", samples=1))
    output_dir = run_dir(root, item)
    write_capture(output_dir / "upstream_token_ids.json", 1, 127)

    records = list(RDLMAdapter().convert_outputs(item, output_dir))
    metadata = json.loads((output_dir / "conversion_metadata.json").read_text())

    assert records[0].token_ids == [11] + [0] * 127
    assert metadata["token_ids_source"] == "retokenized"
    assert metadata["token_ids_transformation"]["reason"] == "upstream_ids_not_canonical_length"


def test_rdlm_retokenizes_out_of_vocab_upstream_rows_offline_and_records_why(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Catch RDLM base-vocabulary IDs being validated as canonical BERT IDs."""

    class FakeTokenizer:
        eos_token = "[SEP]"
        pad_token = "[PAD]"
        pad_token_id = 0

        def __call__(self, texts, **settings):
            assert settings["max_length"] == 128
            return {"input_ids": [[101, 102] + [0] * 126 for _ in texts]}

    class FakeAutoTokenizer:
        @staticmethod
        def from_pretrained(name, **settings):
            assert name == "bert-base-uncased"
            assert settings == {
                "revision": "86b5e0934494bd15c9632b12f734a8a67f723594",
                "local_files_only": True,
            }
            return FakeTokenizer()

    monkeypatch.setitem(
        sys.modules, "transformers", SimpleNamespace(AutoTokenizer=FakeAutoTokenizer)
    )
    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("rdlm", "lm1b", samples=1))
    output_dir = run_dir(root, item)
    write_capture(output_dir / "upstream_token_ids.json", 1, 128)
    capture = json.loads((output_dir / "upstream_token_ids.json").read_text())
    capture["samples"][0]["token_ids"] = [35008] + [1] * 127
    (output_dir / "upstream_token_ids.json").write_text(json.dumps(capture))

    records = list(RDLMAdapter().convert_outputs(item, output_dir))
    metadata = json.loads((output_dir / "conversion_metadata.json").read_text())

    assert records[0].token_ids == [101, 102] + [0] * 126
    assert metadata["token_ids_source"] == "retokenized"
    assert metadata["token_ids_transformation"]["reason"] == "upstream_ids_outside_canonical_vocab"


def test_rdlm_conditional_retokenizes_suffix_out_of_vocab_and_preserves_prefix(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Catch RDLM conditional capture rejecting base-vocabulary suffix IDs."""

    class FakeTokenizer:
        eos_token = "[SEP]"
        pad_token = "[PAD]"
        pad_token_id = 0

        def __call__(self, texts, **settings):
            assert texts == ["generated continuation"]
            assert settings["max_length"] == 64
            return {"input_ids": [[22] * 64]}

    class FakeAutoTokenizer:
        @staticmethod
        def from_pretrained(name, **settings):
            assert name == "bert-base-uncased"
            assert settings == {
                "revision": "86b5e0934494bd15c9632b12f734a8a67f723594",
                "local_files_only": True,
            }
            return FakeTokenizer()

    monkeypatch.setitem(
        sys.modules, "transformers", SimpleNamespace(AutoTokenizer=FakeAutoTokenizer)
    )
    root = prepare_conversion_root(tmp_path)
    item = canonical_conversion_request(root, request("rdlm", "lm1b", steps=1000, samples=1))
    item = replace(
        item,
        generation_mode="conditional_prefix",
        prefix_length=64,
        evaluation_continuation_length=64,
        prompt_count=1024,
        diversity_prompt_count=256,
        completions_per_diversity_prompt=5,
        completion_schedule="quality_then_diversity_v1",
    )
    output_dir = run_dir(root, item)
    prefix = [1] * 64
    suffix = [35008] + [2] * 63
    record = {
        "sample_id": 0,
        "text": "fixed prefix generated continuation",
        "token_ids": prefix + suffix,
        "prompt_id": 0,
        "completion_id": 0,
        "source_index": 100,
        "prefix_token_ids": prefix,
        "reference_token_ids": [3] * 64,
        "full_token_ids": prefix + suffix,
        "prefix_text": "fixed prefix",
        "continuation_text": "generated continuation",
        "reference_text": "source continuation",
        "full_text": "fixed prefix generated continuation",
        "prefix_exact_match": True,
    }
    write_capture(output_dir / "upstream_token_ids.json", 1, 128)
    (output_dir / "upstream_token_ids.json").write_text(
        json.dumps({"schema": "dlb-upstream-token-capture-v1", "samples": [record]}),
        encoding="utf-8",
    )

    records = list(RDLMAdapter().convert_outputs(item, output_dir))
    metadata = json.loads((output_dir / "conversion_metadata.json").read_text())

    assert records[0].prefix_token_ids == prefix
    assert records[0].continuation_token_ids == [22] * 64
    assert records[0].full_token_ids == prefix + [22] * 64
    assert records[0].prefix_exact_match is True
    assert metadata["token_ids_source"] == "retokenized_conditional_suffix"
    assert metadata["token_ids_transformation"]["reason"] == (
        "upstream_ids_outside_canonical_vocab"
    )


def test_shared_capture_runs_langflow_with_locked_offline_tokenizer_and_records_ids(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Catch LangFlow resolving its movable GPT-2 alias or parsing its text file."""

    data_config, downloads, tokenizer_snapshot = prepare_tokenizer_binding(tmp_path, "owt")
    monkeypatch.delenv("HF_HUB_OFFLINE", raising=False)
    monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False)
    entrypoint = tmp_path / "inference.py"
    (entrypoint.parent / "configs").mkdir()
    source = """
import argparse
import os

class Rows:
    def __init__(self, rows): self.rows = rows
    def detach(self): return self
    def cpu(self): return self
    def tolist(self): return self.rows

class LangFlow:
    def generate_samples(self, num_samples, **kwargs):
        return Rows([[index + 1, index + 2] for index in range(num_samples)])

class Tokenizer:
    def batch_decode(self, rows, **kwargs):
        return [f\"text {index}\" for index, _ in enumerate(rows.rows)]

class AutoTokenizer:
    @staticmethod
    def from_pretrained(name, **settings):
        assert name == __SNAPSHOT__
        assert settings == {'local_files_only': True}
        assert os.environ['HF_HUB_OFFLINE'] == '1'
        assert os.environ['TRANSFORMERS_OFFLINE'] == '1'
        return Tokenizer()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_samples', type=int, required=True)
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    rows = LangFlow().generate_samples(args.num_samples)
    texts = AutoTokenizer.from_pretrained('fixture').batch_decode(rows)
    with open(args.output, 'w', encoding='utf-8') as handle:
        handle.write('\\n'.join(texts))
"""
    entrypoint.write_text(
        source.replace("__SNAPSHOT__", repr(str(tokenizer_snapshot))), encoding="utf-8"
    )
    capture_path = tmp_path / "capture.json"
    output_path = tmp_path / "samples.txt"

    result = capture_module.main(
        [
            f"--upstream-entrypoint={entrypoint}",
            f"--capture-path={capture_path}",
            "--capture-kind=langflow",
            "--expected-samples=2",
            f"--data-config-path={data_config}",
            f"--downloads-manifest-path={downloads}",
            "--dataset-id=owt",
            f"--tokenizer-snapshot={tokenizer_snapshot}",
            "--",
            "--num_samples",
            "2",
            "--output",
            str(output_path),
        ]
    )
    capture = json.loads(capture_path.read_text())

    assert result == 0
    assert "HF_HUB_OFFLINE" not in os.environ
    assert "TRANSFORMERS_OFFLINE" not in os.environ
    assert output_path.read_text() == "text 0\ntext 1"
    assert capture == {
        "schema": "dlb-upstream-token-capture-v1",
        "samples": [
            {"sample_id": 0, "text": "text 0", "token_ids": [1, 2]},
            {"sample_id": 1, "text": "text 1", "token_ids": [2, 3]},
        ],
    }


def test_rdlm_capture_routes_saved_assets_and_uses_bounded_exact_batches(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    """Catch movable tokenizers, fixed batches, or fallback assets replacing locked inputs."""

    data_config, downloads, tokenizer_snapshot = prepare_tokenizer_binding(tmp_path, "lm1b")
    monkeypatch.delenv("HF_HUB_OFFLINE", raising=False)
    monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising=False)
    saved_config = tmp_path / "config.yaml"
    saved_config.write_text("model:\n  length: 128\ndata:\n  train: lm1b\n", encoding="utf-8")
    saved_sde = tmp_path / "sde.pkl"
    saved_sde.write_bytes(b"saved-sde")
    checkpoint = tmp_path / "checkpoint.pth"
    checkpoint.write_bytes(b"checkpoint")
    capture_path = tmp_path / "capture.json"

    class Rows:
        def __init__(self, rows):
            self.rows = rows

        def detach(self):
            return self

        def cpu(self):
            return self

        def tolist(self):
            return self.rows

    class FakeTorch:
        @staticmethod
        def load(path, **kwargs):
            assert Path(path) == checkpoint
            return {"config": "embedded-config"}

    class FakeSeqUtils:
        @staticmethod
        def find_bos_and_shift_fn(*args, **kwargs):
            def shift(rows):
                return ["one", "two", "excess"], [
                    [1] * 128,
                    [2] * 128,
                    [3] * 128,
                ]

            return shift

    class FakeBertTokenizer:
        @staticmethod
        def from_pretrained(name, **settings):
            assert name == str(tokenizer_snapshot)
            assert settings == {"local_files_only": True}
            assert os.environ["HF_HUB_OFFLINE"] == "1"
            assert os.environ["TRANSFORMERS_OFFLINE"] == "1"
            return object()

    fake_data = SimpleNamespace(
        transformers=SimpleNamespace(BertTokenizer=FakeBertTokenizer),
        AutoTokenizer=FakeBertTokenizer,
    )

    def get_tokenizer(dataset):
        assert dataset == "lm1b"
        return fake_data.transformers.BertTokenizer.from_pretrained(
            "bert-base-uncased"
        )

    fake_data.get_tokenizer = get_tokenizer

    fake_run_sample = SimpleNamespace(
        torch=FakeTorch(),
        sutils=FakeSeqUtils(),
        data=fake_data,
        tqdm=lambda iterable, **kwargs: iterable,
        instantiate=lambda config, **kwargs: (config, kwargs),
    )
    monkeypatch.setitem(sys.modules, "run_sample", fake_run_sample)

    class FakeOmegaConf:
        @staticmethod
        def load(path):
            assert Path(path) == saved_config
            return SimpleNamespace(
                model=SimpleNamespace(length=128),
                data=SimpleNamespace(train="lm1b"),
            )

        @staticmethod
        def select(config, key):
            section, field = key.split(".")
            return getattr(getattr(config, section), field)

    monkeypatch.setitem(
        sys.modules, "omegaconf", SimpleNamespace(OmegaConf=FakeOmegaConf)
    )

    class FakeMP:
        @staticmethod
        def spawn(*args, **kwargs):
            raise AssertionError("capture must replace process spawning")

    fake_entrypoint = SimpleNamespace(mp=FakeMP())

    def fake_main():
        import run_sample

        def worker(rank, marker):
            assert rank == 0 and marker == "worker"
            run_sample.data.get_tokenizer("lm1b")
            state = run_sample.torch.load(checkpoint)
            assert state["config"].model.length == 128
            with run_sample.open(tmp_path / "wrong" / "sde.pkl", "rb") as handle:
                assert handle.read() == b"saved-sde"
            with pytest.raises(ValueError, match="saved SDE"):
                run_sample.instantiate(
                    "sde",
                    manifold=object(),
                    scheduler=object(),
                    prior_dist=object(),
                    device="cpu",
                )
            run_sample.instantiate(
                "sde",
                manifold=object(),
                scheduler=object(),
                prior_dist=object(),
                preprocessed=(object(), object()),
            )
            assert list(run_sample.tqdm(range(8), leave=False)) == [0, 1]
            shift = run_sample.sutils.find_bos_and_shift_fn(3, 10, object())
            shift(Rows([[0], [0]]))

        fake_entrypoint.mp.spawn(worker, args=("worker",), nprocs=1, join=True)

    fake_entrypoint.main = fake_main
    invocation = capture_module.CaptureInvocation(
        entrypoint=tmp_path / "main.py",
        capture_path=capture_path,
        kind="rdlm",
        expected_samples=2,
        saved_config_path=saved_config,
        saved_sde_path=saved_sde,
        data_config_path=data_config,
        downloads_manifest_path=downloads,
        dataset_id="lm1b",
        tokenizer_snapshot=tokenizer_snapshot,
    )

    capture_module._capture_rdlm(
        fake_entrypoint,
        invocation,
        [f"model_path={checkpoint}", "sampling.batch_per_gpu=1"],
    )
    capture = json.loads(capture_path.read_text())

    assert [sample["text"] for sample in capture["samples"]] == ["one", "two"]
    assert capture["samples"][1]["token_ids"] == [2] * 128
    assert "HF_HUB_OFFLINE" not in os.environ
    assert "TRANSFORMERS_OFFLINE" not in os.environ


def test_continuous_cli_dry_run_emits_valid_commands_and_step_override_rejection(
    tmp_path: Path,
    capsys: pytest.CaptureFixture[str],
) -> None:
    """Catch unsupported cells being substituted or dry-run trying to write artifacts."""

    root = prepare_command_root(tmp_path)
    result = command_main(
        [
            "--root",
            str(root),
            "--models",
            "langflow,rdlm",
            "--datasets",
            "lm1b,owt",
            "--dry-run",
        ]
    )
    records = [json.loads(line) for line in capsys.readouterr().out.splitlines()]

    assert result == 1
    assert len(records) == 4
    supported = [record for record in records if record["status"] == "supported"]
    unsupported = [record for record in records if record["status"] == "unsupported"]
    errors = [record for record in records if record["status"] == "error"]
    assert {(record["model"], record["dataset"]) for record in supported} == {
        ("langflow", "lm1b"),
        ("langflow", "owt"),
    }
    assert {(record["model"], record["dataset"]) for record in unsupported} == {
        ("rdlm", "owt"),
    }
    assert [(record["model"], record["dataset"]) for record in errors] == [
        ("rdlm", "lm1b")
    ]
    assert "allowed: 1000,1024" in errors[0]["reason"]
    assert all(record["command"] for record in supported)
    assert all(record["reason"] for record in unsupported)
    assert all(
        override(record["command"], "seed") == "42"
        if record["model"] == "rdlm"
        else option(record["command"], "--seed") == "42"
        for record in supported
    )
    rdlm_result = command_main(
        [
            "--root",
            str(root),
            "--models",
            "rdlm",
            "--datasets",
            "lm1b",
            "--steps",
            "1000",
            "--dry-run",
        ]
    )
    assert rdlm_result == 0
    rdlm_record = json.loads(capsys.readouterr().out.strip())
    assert rdlm_record["status"] == "supported"
    assert override(rdlm_record["command"], "sampling.batch_per_gpu") == "8"