social-data-attribution-w2/tests/test_draw_working_sample_cli.py

"""Tests for the draw_working_sample CLI entry point."""

from __future__ import annotations

import json

import pandas as pd
import pytest

from data_attribution.recipes.draw_working_sample import _parse_args, main
from dolma.constants import (
    WORKING_SAMPLE_MIN_TOKEN_COUNT,
    WORKING_SAMPLE_SAMPLING_SEED,
)


class TestParseArgs:
    def test_defaults(self):
        args = _parse_args(["--dummy"])
        assert args.dummy is True
        assert args.manifest is None
        assert args.seed == WORKING_SAMPLE_SAMPLING_SEED
        assert args.min_token_count is None
        assert args.max_token_count is None
        assert args.token_floor_per_bin is None
        assert args.docs_per_bin is None

    def test_manifest_path(self, tmp_path):
        args = _parse_args(["--manifest", str(tmp_path / "m.parquet")])
        assert args.manifest == tmp_path / "m.parquet"

    def test_min_max_token_count(self):
        args = _parse_args(
            ["--dummy", "--min-token-count", "256", "--max-token-count", "4096"]
        )
        assert args.min_token_count == 256
        assert args.max_token_count == 4096

    def test_docs_per_bin_and_token_floor_mutually_exclusive(self):
        with pytest.raises(SystemExit):
            _parse_args(
                ["--dummy", "--docs-per-bin", "10", "--token-floor-per-bin", "1000"]
            )

    def test_seed_override(self):
        args = _parse_args(["--dummy", "--seed", "99"])
        assert args.seed == 99


class TestMainDummyMode:
    def test_dummy_produces_outputs(self, tmp_path):
        result = main(
            [
                "--dummy",
                "--n-dummy-docs",
                "5000",
                "--output-dir",
                str(tmp_path),
                "--docs-per-bin",
                "5",
            ]
        )
        assert result == 0
        assert (tmp_path / "working_sample_manifest.parquet").exists()
        assert (tmp_path / "sample_contract.json").exists()
        assert (tmp_path / "bin_summary.csv").exists()

    def test_dummy_applies_default_min_token_filter(self, tmp_path):
        result = main(
            [
                "--dummy",
                "--n-dummy-docs",
                "10000",
                "--output-dir",
                str(tmp_path),
                "--docs-per-bin",
                "5",
            ]
        )
        assert result == 0
        contract = json.loads((tmp_path / "sample_contract.json").read_text())
        assert (
            contract["WORKING_SAMPLE_MIN_TOKEN_COUNT"] == WORKING_SAMPLE_MIN_TOKEN_COUNT
        )

        manifest = pd.read_parquet(tmp_path / "working_sample_manifest.parquet")
        assert (manifest["token_count"] >= WORKING_SAMPLE_MIN_TOKEN_COUNT).all()

    def test_dummy_respects_explicit_min_token(self, tmp_path):
        result = main(
            [
                "--dummy",
                "--n-dummy-docs",
                "10000",
                "--output-dir",
                str(tmp_path),
                "--docs-per-bin",
                "5",
                "--min-token-count",
                "1024",
            ]
        )
        assert result == 0
        contract = json.loads((tmp_path / "sample_contract.json").read_text())
        assert contract["WORKING_SAMPLE_MIN_TOKEN_COUNT"] == 1024

        manifest = pd.read_parquet(tmp_path / "working_sample_manifest.parquet")
        assert (manifest["token_count"] >= 1024).all()

    def test_no_manifest_no_dummy_returns_error(self, tmp_path):
        result = main(["--output-dir", str(tmp_path)])
        assert result == 1


class TestExcludeManifestCli:
    def test_parse_exclude_manifest(self, tmp_path):
        p = tmp_path / "exclude.parquet"
        p.touch()
        args = _parse_args(["--dummy", "--exclude-manifest", str(p)])
        assert args.exclude_manifest == [p]

    def test_parse_multiple_exclude_manifests(self, tmp_path):
        p1 = tmp_path / "e1.parquet"
        p2 = tmp_path / "e2.parquet"
        p1.touch()
        p2.touch()
        args = _parse_args(
            [
                "--dummy",
                "--exclude-manifest",
                str(p1),
                "--exclude-manifest",
                str(p2),
            ]
        )
        assert args.exclude_manifest == [p1, p2]

    def test_exclude_manifest_removes_docs(self, tmp_path):
        out_a = tmp_path / "sample_a"
        result_a = main(
            [
                "--dummy",
                "--n-dummy-docs",
                "5000",
                "--output-dir",
                str(out_a),
                "--docs-per-bin",
                "3",
            ]
        )
        assert result_a == 0
        manifest_a = pd.read_parquet(out_a / "working_sample_manifest.parquet")
        ids_a = set(manifest_a["doc_id"])

        out_b = tmp_path / "sample_b"
        result_b = main(
            [
                "--dummy",
                "--n-dummy-docs",
                "5000",
                "--output-dir",
                str(out_b),
                "--docs-per-bin",
                "3",
                "--exclude-manifest",
                str(out_a / "working_sample_manifest.parquet"),
            ]
        )
        assert result_b == 0
        manifest_b = pd.read_parquet(out_b / "working_sample_manifest.parquet")
        ids_b = set(manifest_b["doc_id"])

        assert ids_a.isdisjoint(ids_b)

    def test_exclude_manifest_contract_recorded(self, tmp_path):
        out_a = tmp_path / "sample_a"
        main(
            [
                "--dummy",
                "--n-dummy-docs",
                "5000",
                "--output-dir",
                str(out_a),
                "--docs-per-bin",
                "3",
            ]
        )

        out_b = tmp_path / "sample_b"
        exclude_path = out_a / "working_sample_manifest.parquet"
        main(
            [
                "--dummy",
                "--n-dummy-docs",
                "5000",
                "--output-dir",
                str(out_b),
                "--docs-per-bin",
                "3",
                "--exclude-manifest",
                str(exclude_path),
            ]
        )

        contract = json.loads((out_b / "sample_contract.json").read_text())
        assert "WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS" in contract
        assert contract["WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS"] == [str(exclude_path)]
        assert contract["WORKING_SAMPLE_EXCLUDED_DOC_COUNT"] > 0


class TestEndToEndSamplingPipeline:
    def test_load_filter_draw_write(self, tmp_path):
        from dolma.paper.sampling_loader import apply_token_filter
        from dolma.pool_sample.sampling import generate_dummy_manifest
        from dolma.working_sample import (
            draw_working_sample,
            sample_contract,
            write_outputs,
        )

        manifest = generate_dummy_manifest(n_docs=10_000, seed=42)
        original_count = len(manifest)
        filtered = apply_token_filter(manifest, "token_count", min_tokens=512)
        assert len(filtered) < original_count
        assert (filtered["token_count"] >= 512).all()

        result = draw_working_sample(
            filtered,
            docs_per_bin=5,
            seed=42,
            min_token_count=512,
        )
        assert len(result.sample_df) > 0
        assert (result.sample_df["token_count"] >= 512).all()

        write_outputs(result, tmp_path)
        contract = sample_contract(result)
        assert contract["WORKING_SAMPLE_MIN_TOKEN_COUNT"] == 512
        assert contract["WORKING_SAMPLE_MAX_TOKEN_COUNT"] is None
        assert contract["WORKING_SAMPLE_REALIZED_DOC_COUNT"] == len(result.sample_df)

        written_manifest = pd.read_parquet(tmp_path / "working_sample_manifest.parquet")
        assert len(written_manifest) == len(result.sample_df)