[test] add test for cli

tests/cli/conftest.py

@@ -0,0 +1,138 @@
+"""Shared fixtures for end-to-end CLI tests.
+
+Each test invokes the Typer app in-process via ``CliRunner`` and compares the
+captured stdout (never stderr) against an exact snapshot. All env vars bound
+to CLI options are cleared, and ``COLUMNS``/``NO_COLOR`` are pinned so that
+rich renders logs at a fixed width without ANSI codes.
+"""
+
+import json
+from fnmatch import fnmatch
+
+import pytest
+from typer.testing import CliRunner
+
+# env vars bound to typer options in prep.cli; keep in sync with cli.py
+CLI_ENVVARS = (
+    "SRC",
+    "HEAD",
+    "TAIL",
+    "NPROC",
+    "SEED",
+    "SAVE",
+    "SAVE_DIR",
+    "SAVE_PARQ",
+    "SAVE_NPROC",
+    "SAVE_PREVIEW",
+    "HF",
+    "HF_REPO",
+    "HF_SUBSET",
+    "HF_PRIVATE",
+    "HF_NPROC",
+    "MAX_SAMPLES",
+    "Q_COLS",
+    "Q_TEMP",
+    "OP_COLS",
+    "A_COLS",
+    "A_TEMP",
+    "VERL_ABILITY",
+    "VERL_STYLE",
+    "LABELS",
+    "TRANSFORMS",
+    "UI",
+    "COLUMNS",
+    "NO_COLOR",
+    "FORCE_COLOR",
+)
+
+
+@pytest.fixture(autouse=True)
+def cli_env(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, tmp_path):
+    """Isolate the CLI from the host environment and fix rich rendering."""
+    import logging
+
+    from rich.logging import RichHandler
+
+    for var in CLI_ENVVARS:
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("COLUMNS", "120")
+    monkeypatch.setenv("NO_COLOR", "1")
+    # RichHandler captures an os.environ snapshot at import time, while pytest
+    # capture swaps os.environ per test; pin the cached console size directly
+    # so log records always render at a fixed width.
+    for handler in logging.getLogger().handlers:
+        if isinstance(handler, RichHandler):
+            console = handler.console
+            prev_size = (console._width, console._height)
+            console.size = (120, 25)
+
+            def restore(prev=prev_size, c=console):
+                c._width, c._height = prev
+
+            request.addfinalizer(restore)
+    monkeypatch.chdir(tmp_path)
+
+
+@pytest.fixture
+def runner() -> CliRunner:
+    """CliRunner for invoking the prep/ppls commands in-process."""
+    return CliRunner()
+
+
+@pytest.fixture
+def qa_jsonl(tmp_path) -> str:
+    """A 10-row question/answer source in jsonl format."""
+    path = tmp_path / "qa.jsonl"
+    path.write_text(
+        "\n".join(
+            json.dumps({"question": f"What is {i}+{i}?", "answer": str(2 * i)})
+            for i in range(10)
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    return path.name
+
+
+@pytest.fixture
+def mc_jsonl(tmp_path) -> str:
+    """A 10-row multiple-choice source in jsonl format."""
+    path = tmp_path / "mc.jsonl"
+    path.write_text(
+        "\n".join(
+            json.dumps(
+                {
+                    "question": f"Pick the even number among the choices (set {i}).",
+                    "options": [str(2 * i + 1), str(2 * i + 2), str(2 * i + 3)],
+                    "answer": "B",
+                }
+            )
+            for i in range(10)
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+    return path.name
+
+
+def _rstrip_lines(text: str) -> list[str]:
+    return [line.rstrip() for line in text.splitlines()]
+
+
+@pytest.fixture
+def assert_stdout():
+    """Assert exact stdout equality (modulo rich's right-padding of log lines)."""
+
+    def _assert(result, expected: str, exit_code: int = 0) -> None:
+        assert result.exit_code == exit_code
+        for line, exp in zip(
+            _rstrip_lines(result.stdout), _rstrip_lines(expected), strict=True
+        ):
+            if "*" in exp:
+                assert fnmatch(line, exp), (
+                    f"line '{line}' does not match pattern '{exp}'"
+                )
+            else:
+                assert line == exp
+
+    return _assert

tests/cli/test_clip.py

@@ -0,0 +1,20 @@
+"""CLI tests for the clip target format (exact stdout comparison).
+
+No pipelines are registered for 'clip', so the CLI falls back to 'auto',
+which is also unregistered for this format, and fails with a KeyError.
+"""
+
+from prep.cli import app
+
+EXPECTED = """\
+WARNING  ⚠️      Formatter pipeline not registered for ('auto', 'clip', 'train').
+WARNING  Fallback to generic pipeline 'auto', which may cause unexpected formatting issues."""
+
+
+def test_prep_clip_fallback_fails(runner, qa_jsonl, assert_stdout):
+    result = runner.invoke(
+        app, ["prep", "clip", "auto", "train", qa_jsonl, "--nproc", "1"]
+    )
+
+    assert_stdout(result, EXPECTED, exit_code=1)
+    assert isinstance(result.exception, KeyError)

tests/cli/test_cls.py

@@ -0,0 +1,36 @@
+import pytest
+from typer.testing import CliRunner
+
+from prep.cli import app
+
+EXPECTED = """\
+WARNING  ⚠️      Formatter pipeline not registered for ('cifar10', 'cls', 'train').
+WARNING  Fallback to generic pipeline 'auto', which may cause unexpected formatting issues.
+INFO     Dataset({
+             features: ['id', 'image', 'label', 'extra_info'],
+             num_rows: 50000
+         })
+INFO     {'id': 'cifar10/train00000000', 'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=32x32 at
+         *>, 'label': 0, 'extra_info': ''}
+💾	About to save to disk -> PosixPath('out/cls/cifar10/train')
+💾	Skipping (non-interactive mode)
+☁️\tAbout to upload to HF -> 'cifar10' (subset='default', split='train', private=True)
+☁️\tSkipping (non-interactive mode)"""
+
+
+@pytest.mark.slow
+def test_prep_cls_stdout(runner: CliRunner, assert_stdout):
+    result = runner.invoke(
+        app,
+        [
+            "prep",
+            "cls",
+            "cifar10",
+            "train",
+            "uoft-cs/cifar10",
+            "--head",
+            "1",
+        ],
+    )
+
+    assert_stdout(result, EXPECTED)

tests/cli/test_eval.py

@@ -0,0 +1,23 @@
+"""CLI tests for the eval target format (exact stdout comparison)."""
+
+from prep.cli import app
+
+EXPECTED = """\
+INFO     Dataset({
+             features: ['id', 'images', 'question', 'options', 'answer'],
+             num_rows: 10
+         })
+INFO     {'id': 'mc.jsonl/test00000000', 'images': [], 'question': 'Pick the even number among the choices (set 0).',
+         'options': ['1', '2', '3'], 'answer': 'B'}
+💾	About to save to disk -> PosixPath('out/eval/auto/test')
+💾	Skipping (non-interactive mode)
+☁️	About to upload to HF -> 'auto' (subset='default', split='test', private=True)
+☁️	Skipping (non-interactive mode)"""
+
+
+def test_prep_eval_stdout(runner, mc_jsonl, assert_stdout):
+    result = runner.invoke(
+        app, ["prep", "eval", "auto", "test", mc_jsonl, "--head", "1", "--nproc", "1"]
+    )
+
+    assert_stdout(result, EXPECTED)

tests/cli/test_ppls.py

@@ -0,0 +1,53 @@
+"""CLI tests for the ppls command (exact stdout comparison)."""
+
+from prep.cli import app
+
+LIST_FORMAT_EXPECTED = """\
+sft
+verl
+eval
+clip
+cls
+others"""
+
+TABLE_EXPECTED = """\
+Showing results under 'out'
+                                  sft
+┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━┳━━━━━━━┳━━━━━┳━━━━━━┓
+┃ ID    ┃ Default Source     ┃ Local Path ┃ Size ┃ train ┃ val ┃ test ┃
+┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━╇━━━━━━━╇━━━━━╇━━━━━━┩
+│ auto  │                    │            │      │ ✖     │ ✖   │ ✖    │
+│ geo3k │ hiyouga/geometry3k │            │      │ ✖     │     │ ✖    │
+└───────┴────────────────────┴────────────┴──────┴───────┴─────┴──────┘
+                                 verl
+┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━┳━━━━━━━┳━━━━━┳━━━━━━┓
+┃ ID    ┃ Default Source     ┃ Local Path ┃ Size ┃ train ┃ val ┃ test ┃
+┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━╇━━━━━━━╇━━━━━╇━━━━━━┩
+│ auto  │                    │            │      │ ✖     │ ✖   │ ✖    │
+│ geo3k │ hiyouga/geometry3k │            │      │ ✖     │     │ ✖    │
+└───────┴────────────────────┴────────────┴──────┴───────┴─────┴──────┘
+                               eval
+┏━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━┳━━━━━━━┳━━━━━┳━━━━━━┓
+┃ ID   ┃ Default Source ┃ Local Path ┃ Size ┃ train ┃ val ┃ test ┃
+┡━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━╇━━━━━━━╇━━━━━╇━━━━━━┩
+│ auto │                │            │      │ ✖     │ ✖   │ ✖    │
+└──────┴────────────────┴────────────┴──────┴───────┴─────┴──────┘
+
+                               cls
+┏━━━━━━┳━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━┳━━━━━━━┳━━━━━┳━━━━━━┓
+┃ ID   ┃ Default Source ┃ Local Path ┃ Size ┃ train ┃ val ┃ test ┃
+┡━━━━━━╇━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━╇━━━━━━━╇━━━━━╇━━━━━━┩
+│ auto │                │            │      │ ✖     │ ✖   │ ✖    │
+└──────┴────────────────┴────────────┴──────┴───────┴─────┴──────┘"""
+
+
+def test_ppls_list_format_stdout(runner, assert_stdout):
+    result = runner.invoke(app, ["ppls", "--list-format"])
+
+    assert_stdout(result, LIST_FORMAT_EXPECTED)
+
+
+def test_ppls_table_stdout(runner, assert_stdout):
+    result = runner.invoke(app, ["ppls"])
+
+    assert_stdout(result, TABLE_EXPECTED)

tests/cli/test_sft.py

@@ -0,0 +1,23 @@
+"""CLI tests for the sft target format (exact stdout comparison)."""
+
+from prep.cli import app
+
+EXPECTED = """\
+INFO     Dataset({
+             features: ['images', 'messages', 'id', 'extra_info'],
+             num_rows: 10
+         })
+INFO     {'images': [], 'messages': [{'role': 'user', 'content': 'What is 0+0?'}, {'role': 'assistant', 'content':
+         '0'}], 'id': 'qa.jsonl/00000000', 'extra_info': ''}
+💾	About to save to disk -> PosixPath('out/sft/auto/train')
+💾	Skipping (non-interactive mode)
+☁️	About to upload to HF -> 'auto' (subset='default', split='train', private=True)
+☁️	Skipping (non-interactive mode)"""
+
+
+def test_prep_sft_stdout(runner, qa_jsonl, assert_stdout):
+    result = runner.invoke(
+        app, ["prep", "sft", "auto", "train", qa_jsonl, "--head", "1", "--nproc", "1"]
+    )
+
+    assert_stdout(result, EXPECTED)

tests/cli/test_show.py

@@ -0,0 +1,22 @@
+"""CLI tests for the show target format (exact stdout comparison)."""
+
+from prep.cli import app
+
+EXPECTED = """\
+INFO     Dataset({
+             features: ['question', 'answer'],
+             num_rows: 10
+         })
+INFO     {'question': 'What is 0+0?', 'answer': '0'}
+💾	About to save to disk -> PosixPath('out/show/_/train')
+💾	Skipping (non-interactive mode)
+☁️	About to upload to HF -> '_' (subset='default', split='train', private=True)
+☁️	Skipping (non-interactive mode)"""
+
+
+def test_prep_show_stdout(runner, qa_jsonl, assert_stdout):
+    result = runner.invoke(
+        app, ["prep", "show", "_", "train", qa_jsonl, "--head", "1", "--nproc", "1"]
+    )
+
+    assert_stdout(result, EXPECTED)

tests/cli/test_verl.py

@@ -0,0 +1,38 @@
+"""CLI tests for the verl target format (exact stdout comparison)."""
+
+from prep.cli import app
+
+EXPECTED = """\
+INFO     Dataset({
+             features: ['images', 'data_source', 'prompt', 'ability', 'reward_model', 'extra_info'],
+             num_rows: 10
+         })
+INFO     {'images': [], 'data_source': 'qa.jsonl', 'prompt': [{'role': 'user', 'content': 'What is 0+0?\\nPlease reason
+         step by step, and put your final answer within \\\\boxed{}.'}], 'ability': 'math', 'reward_model': {'style':
+         'rule', 'ground_truth': '0'}, 'extra_info': {'split': 'train', 'index': '00000000', 'explanation': '', 'misc':
+         ''}}
+💾	About to save to disk -> PosixPath('out/verl/auto/train')
+💾	Skipping (non-interactive mode)
+☁️	About to upload to HF -> 'auto' (subset='default', split='train', private=True)
+☁️	Skipping (non-interactive mode)"""
+
+
+def test_prep_verl_stdout(runner, qa_jsonl, assert_stdout):
+    result = runner.invoke(
+        app,
+        [
+            "prep",
+            "verl",
+            "auto",
+            "train",
+            qa_jsonl,
+            "--head",
+            "1",
+            "--nproc",
+            "1",
+            "--q-template",
+            "{question}\nPlease reason step by step, and put your final answer within \\boxed{{}}.",
+        ],
+    )
+
+    assert_stdout(result, EXPECTED)