social-data-attribution-w2/tests/data_attribution/test_cli_defaults.py

import os
from pathlib import Path
from types import SimpleNamespace

import pytest

from data_attribution.cli.handlers import _attribution_handler
from data_attribution import config
from data_attribution.recipes import dolma_training_index, socialiqa_queries


def test_dolma_training_index_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("TRAIN_SLICE", raising=False)
    monkeypatch.delenv("MODEL_ID", raising=False)
    args = dolma_training_index._parse_args([])
    assert args.train_slice == 4509
    assert args.model_id == "allenai/Olmo-3-1025-7B"


def test_socialiqa_queries_defaults(monkeypatch: pytest.MonkeyPatch) -> None:
    for var in ("SPLIT", "SAMPLE_SIZE", "MODEL_ID"):
        monkeypatch.delenv(var, raising=False)
    args = socialiqa_queries._parse_args([])
    assert args.split == "train"
    assert args.sample_size == 100
    assert args.model_id == "allenai/Olmo-3-1025-7B"


def test_hf_cache_default_path(monkeypatch: pytest.MonkeyPatch) -> None:
    for var in ("HF_CACHE_ROOT", "HF_HOME"):
        monkeypatch.delenv(var, raising=False)
    path = config.resolve_hf_cache_root()
    expected = config.PROJECT_ROOT / ".hf_cache"
    assert path == expected


def test_configure_hf_cache_clears_deprecated_transformers_env(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    cache_root = tmp_path / "hf_cache"
    for var in (
        "HF_CACHE_ROOT",
        "HF_HOME",
        "HUGGINGFACE_HUB_CACHE",
        "HF_DATASETS_CACHE",
        *config.DEPRECATED_TRANSFORMERS_CACHE_VARS,
    ):
        monkeypatch.delenv(var, raising=False)
    for var in config.DEPRECATED_TRANSFORMERS_CACHE_VARS:
        monkeypatch.setenv(var, str(tmp_path / f"stale_{var.lower()}"))

    paths = config.configure_hf_cache(cache_root, create_dirs=True)

    assert os.environ["HF_HOME"] == str(paths.root)
    assert os.environ["HUGGINGFACE_HUB_CACHE"] == str(paths.hub)
    assert os.environ["HF_DATASETS_CACHE"] == str(paths.datasets)
    assert paths.root == cache_root
    assert paths.hub.is_dir()
    assert paths.datasets.is_dir()
    assert paths.transformers.is_dir()
    for var in config.DEPRECATED_TRANSFORMERS_CACHE_VARS:
        assert var not in os.environ


def test_attribution_handler_sets_default_output_dir(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    captured: dict[str, Path] = {}

    def fake_run_and_store(config):
        captured["output_dir"] = config.output_dir
        return tmp_path / "out.parquet"

    monkeypatch.setattr(
        "data_attribution.attribution.run_and_store", fake_run_and_store
    )

    args = SimpleNamespace(
        command="run",
        run_id=None,
        query_gradients=tmp_path / "queries",
        index_path=tmp_path / "index",
        output_dir=None,
        metadata_path=None,
        output_format="parquet",
        top_k=10,
        query_id_field="query_id",
        doc_id_field="doc_id",
        metadata_join_keys=None,
        device=None,
        unit_norm=True,
        metadata=None,
    )
    result = _attribution_handler(args, {})
    assert result == 0
    assert "output_dir" in captured
    default_dir = captured["output_dir"]
    assert default_dir.parts[:2] == ("runs", "attributions")
    assert default_dir.name.startswith("attr-")