audit-labs/audit-tools

 1from types import SimpleNamespace
 2
 3import pandas as pd
 4
 5from sampling.sampling_tool.cli import run
 6
 7
 8def _write_population(path, rows=20):
 9    frame = pd.DataFrame(
10        {
11            "ID": [f"ID{i:03d}" for i in range(rows)],
12            "Status": ["Closed"] * rows,
13        }
14    )
15    frame.to_csv(path, index=False)
16
17
18def _options(input_path, out_path, seed=123, sample_size=5, method="random", **kwargs):
19    values = {
20        "input": str(input_path),
21        "sheet": None,
22        "id_column": "ID",
23        "method": method,
24        "sample_size": sample_size,
25        "stratify_column": None,
26        "strata_counts": None,
27        "strata_proportions": None,
28        "seed": seed,
29        "out": str(out_path),
30        "exclude_blank_id": False,
31        "dedupe_id": "fail",
32        "filters": {},
33        "allow_shortfall": False,
34    }
35    values.update(kwargs)
36    return SimpleNamespace(**values)
37
38
39def test_random_sample_returns_correct_size(tmp_path):
40    source = tmp_path / "population.csv"
41    _write_population(source)
42
43    run_dir = run(_options(source, tmp_path / "out"))
44
45    sample = pd.read_csv(run_dir / "sample.csv")
46    assert len(sample) == 5
47
48
49def test_same_seed_returns_same_selected_ids(tmp_path):
50    source = tmp_path / "population.csv"
51    _write_population(source)
52
53    first = run(_options(source, tmp_path / "out1", seed=20260707))
54    second = run(_options(source, tmp_path / "out2", seed=20260707))
55
56    first_ids = pd.read_csv(first / "sample.csv")["ID"].tolist()
57    second_ids = pd.read_csv(second / "sample.csv")["ID"].tolist()
58    assert first_ids == second_ids
59
60
61def test_different_seed_can_return_different_selected_ids(tmp_path):
62    source = tmp_path / "population.csv"
63    _write_population(source)
64
65    first = run(_options(source, tmp_path / "out1", seed=1))
66    second = run(_options(source, tmp_path / "out2", seed=2))
67
68    first_ids = pd.read_csv(first / "sample.csv")["ID"].tolist()
69    second_ids = pd.read_csv(second / "sample.csv")["ID"].tolist()
70    assert first_ids != second_ids