Implement census analysis: read, count, summarise, write and CLI

Closes #11 Closes #12 Closes #13 Closes #14 Closes #15 Task: MIL-002#4 Task: MIL-002#5 Task: MIL-002#6 Task: MIL-002#7 Task: MIL-002#8
2026-10-13 15:37:51 +08:00
parent 259d017ad2
commit 7459711243
9 changed files with 364 additions and 0 deletions
@@ -0,0 +1,89 @@
+"""Tests for the squirrel.census module."""
+
+import pathlib
+
+from squirrel import census
+from squirrel.constants import (
+    FILE_ENCODING,
+    FUR_COLOR_COLUMN,
+    FUR_COLORS,
+    SUMMARY_COLUMNS,
+)
+
+OTHER_COLUMN = "Unique Squirrel ID"
+
+
+def _row(color: str) -> dict[str, str]:
+    return {FUR_COLOR_COLUMN: color}
+
+
+def test_read_csv_returns_one_dict_per_row(tmp_path: pathlib.Path) -> None:
+    path = tmp_path / "census.csv"
+    path.write_text(
+        f"{OTHER_COLUMN},{FUR_COLOR_COLUMN}\n1,{FUR_COLORS[0]}\n2,\n",
+        encoding=FILE_ENCODING,
+    )
+
+    rows = census.read_csv(path)
+
+    assert rows == [
+        {OTHER_COLUMN: "1", FUR_COLOR_COLUMN: FUR_COLORS[0]},
+        {OTHER_COLUMN: "2", FUR_COLOR_COLUMN: ""},
+    ]
+
+
+def test_count_by_fur_color_counts_each_colour() -> None:
+    gray, cinnamon, black = FUR_COLORS
+    rows = [
+        _row(gray),
+        _row(gray),
+        _row(cinnamon),
+        _row(black),
+        _row(black),
+        _row(black),
+    ]
+
+    counts = census.count_by_fur_color(rows)
+
+    assert counts == {gray: 2, cinnamon: 1, black: 3}
+
+
+def test_count_by_fur_color_ignores_blank_and_unknown_colours() -> None:
+    rows = [_row(""), _row("White"), _row(FUR_COLORS[0].lower()), {OTHER_COLUMN: "1"}]
+
+    counts = census.count_by_fur_color(rows)
+
+    assert sum(counts.values()) == 0
+
+
+def test_count_by_fur_color_includes_zero_counts_in_colour_order() -> None:
+    counts = census.count_by_fur_color([_row(FUR_COLORS[1])])
+
+    assert list(counts.items()) == [
+        (FUR_COLORS[0], 0),
+        (FUR_COLORS[1], 1),
+        (FUR_COLORS[2], 0),
+    ]
+
+
+def test_count_by_fur_color_uses_given_colours() -> None:
+    counts = census.count_by_fur_color([_row("White"), _row("White")], colors=["White"])
+
+    assert counts == {"White": 2}
+
+
+def test_build_summary_keeps_order_of_counts() -> None:
+    counts = {"Black": 3, "Gray": 1, "Cinnamon": 2}
+
+    assert census.build_summary(counts) == [("Black", 3), ("Gray", 1), ("Cinnamon", 2)]
+
+
+def test_write_csv_writes_header_then_rows(tmp_path: pathlib.Path) -> None:
+    path = tmp_path / "summary.csv"
+    summary = [(FUR_COLORS[0], 5), (FUR_COLORS[1], 0)]
+
+    census.write_csv(summary, path)
+
+    header = ",".join(SUMMARY_COLUMNS)
+    expected = f"{header}\n{FUR_COLORS[0]},5\n{FUR_COLORS[1]},0\n"
+    assert path.read_bytes().decode(FILE_ENCODING).replace("\r\n", "\n") == expected