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):
 9    frame = pd.DataFrame(
10        {
11            "ID": [f"ID{i:03d}" for i in range(12)],
12            "Type": ["A"] * 5 + ["B"] * 4 + ["C"] * 3,
13        }
14    )
15    frame.to_csv(path, index=False)
16
17
18def _options(input_path, out_path, **kwargs):
19    values = {
20        "input": str(input_path),
21        "sheet": None,
22        "id_column": "ID",
23        "method": "stratified",
24        "sample_size": None,
25        "stratify_column": "Type",
26        "strata_counts": "A=2,B=2,C=1",
27        "strata_proportions": None,
28        "seed": 50,
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_stratified_counts_select_exact_requested_counts(tmp_path):
40    source = tmp_path / "population.csv"
41    _write_population(source)
42
43    run_dir = run(_options(source, tmp_path / "out"))
44
45    counts = pd.read_csv(run_dir / "sample.csv")["Type"].value_counts().to_dict()
46    assert counts == {"A": 2, "B": 2, "C": 1}
47
48
49def test_stratified_proportions_use_largest_remainder(tmp_path):
50    source = tmp_path / "population.csv"
51    _write_population(source)
52
53    run_dir = run(
54        _options(
55            source,
56            tmp_path / "out",
57            sample_size=7,
58            strata_counts=None,
59            strata_proportions="A=0.50,B=0.30,C=0.20",
60        )
61    )
62
63    sample = pd.read_csv(run_dir / "sample.csv")
64    counts = sample["Type"].value_counts().to_dict()
65    assert len(sample) == 7
66    assert counts == {"A": 4, "B": 2, "C": 1}