social-data-attribution-w2/tests/data_attribution/cache/test_policy.py

from pathlib import Path
from types import SimpleNamespace

import pytest

from data_attribution.cache import precache
from data_attribution.cli.ingest import _ingest_handler


def test_precache_rejects_unknown_model(tmp_path: Path) -> None:
    with pytest.raises(ValueError, match="Unknown model id"):
        precache.main(
            [
                "--cache-root",
                str(tmp_path),
                "--model-id",
                "unknown/model",
                "--skip-dataset",
            ]
        )


def test_precache_rejects_unknown_dataset(tmp_path: Path) -> None:
    with pytest.raises(ValueError, match="Unknown dataset id"):
        precache.main(
            [
                "--cache-root",
                str(tmp_path),
                "--dataset-repo-id",
                "unknown/dataset",
                "--skip-model",
                "--skip-dataset",
            ]
        )


def test_ingest_requires_local_shards(tmp_path: Path) -> None:
    args = SimpleNamespace(
        command="ingest",
        run_id=None,
        output_path=tmp_path / "out.jsonl",
        manifest_path=tmp_path / "manifest.json",
        subset=None,
        split="train",
        num_documents=1,
        token_budget=None,
        stratify_by=None,
        output_format="jsonl",
        cache_dir=None,
        download_dir=None,
        shuffle_buffer=10_000,
        seed=0,
    )
    with pytest.raises(
        ValueError, match="Ingest requires --cache-dir or --download-dir"
    ):
        _ingest_handler(args, {})