scripts/generate_normalization_benchmark_fixture.py

#!/usr/bin/env python3
"""Regenerate tests/fixtures/normalization_benchmark_v1.json with core cases + N random samples.

Uses a fixed seed so CI expectations are stable. Run from repo root after changing
merge rules if you intentionally want to refresh sampled rows (then review git diff).

  python scripts/generate_normalization_benchmark_fixture.py
"""
from __future__ import annotations

import argparse
import json
import random
from pathlib import Path

from backend.kink_merge import propose_merge_clusters

_REPO = Path(__file__).resolve().parents[1]
_DEFAULT_OUT = _REPO / "tests" / "fixtures" / "normalization_benchmark_v1.json"

# ASCII-only, length >= SINGLE_TOKEN_MERGE_MIN_CHARS (6); prefer 8+ to stay clear of edge.
_WORD_POOL = [
    "moonbeams",
    "starlight",
    "fireworks",
    "cranberry",
    "waterfall",
    "butterfly",
    "dragonfly",
    "nightfall",
    "riverbank",
    "snowflake",
    "sunflower",
    "pineapple",
    "blueberry",
    "strawberry",
    "raspberry",
    "candlewax",
    "bookshelf",
    "doorframe",
    "sidewalks",
    "skylights",
    "clockwork",
    "dreamland",
    "friendship",
    "leadership",
    "membership",
    "partnership",
    "relationship",
    "scholarship",
    "fellowship",
    "citizenship",
    "apprentice",
    "basketball",
    "volleyball",
    "footprints",
    "headlights",
    "spotlights",
    "flashlight",
    "lighthouse",
    "greenhouse",
    "wheelhouse",
    "glasshouse",
    "farmhouses",
    "workhouses",
    "guesthouse",
    "roadhouses",
    "greenfield",
    "battlefield",
    "minefields",
    "cornfields",
    "haystacks",
]


def _core_cases() -> list[dict]:
    return [
        {
            "rows": [
                ["snug_a", "Snuggles", "fetlife_fetish", 10.0],
                ["snug_b", "snuggles", "fetlife_fetish", 200.0],
            ],
            "expect_canonical_id": "snug_b",
            "expect_duplicate_ids": ["snug_a"],
        },
        {
            "rows": [
                ["dup_fmf", "threesome fmf", "fetlife_fetish", 100.0],
                ["canon_fmf", "FMF threesomes", "fetlife_fetish", 500.0],
                ["mfm_other", "MFM threesomes", "fetlife_fetish", 400.0],
            ],
            "expect_canonical_id": "canon_fmf",
            "expect_duplicate_ids": ["dup_fmf"],
            "expect_excluded_ids": ["mfm_other"],
        },
    ]


def _random_case_only_clusters(*, rng: random.Random, n: int) -> list[dict]:
    words = rng.sample(_WORD_POOL, k=min(n, len(_WORD_POOL)))
    cases: list[dict] = []
    for i, w in enumerate(words):
        lo_id = f"rnd_{i:02d}_lo"
        hi_id = f"rnd_{i:02d}_hi"
        cluster = f"sbench_{rng.randint(0, 9999):04d}"
        rows = [
            [lo_id, w.capitalize(), cluster, float(rng.randint(1, 50))],
            [hi_id, w, cluster, float(rng.randint(60, 200))],
        ]
        tuples = [tuple(r) for r in rows]
        clusters = propose_merge_clusters(tuples)
        assert len(clusters) == 1, (w, clusters)
        c0 = clusters[0]
        case: dict = {
            "rows": rows,
            "expect_canonical_id": c0["canonical_id"],
            "expect_duplicate_ids": list(c0["duplicate_ids"]),
        }
        cases.append(case)
    return cases


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--out", type=Path, default=_DEFAULT_OUT, help="fixture path to write")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--random-cases", type=int, default=20, help="number of seeded random case-only clusters")
    args = ap.parse_args()

    rng = random.Random(args.seed)  # noqa: S311 — deterministic benchmark sampling, not crypto
    payload = {
        "comment": (
            "Golden expectations for propose_merge_clusters; extend after reviewing export JSONL. "
            "Random block was generated by scripts/generate_normalization_benchmark_fixture.py."
        ),
        "random_seed": args.seed,
        "random_case_only_count": args.random_cases,
        "merge_cluster_cases": _core_cases() + _random_case_only_clusters(rng=rng, n=args.random_cases),
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(payload, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(json.dumps({"wrote": str(args.out), "cases": len(payload["merge_cluster_cases"])}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())