Northern-Thailand-Ping-River-Monitor/rust/collector/testdata/gen_notify_cases.py

"""Notification cases for rust/collector/src/notify.rs, from src/notify.py itself.

    uv run python rust/collector/testdata/gen_notify_cases.py

Random hourly sequences around the thresholds (capacity guard, hysteresis,
danger escalation, the P.1 outlook, feed gaps, ntfy outages for some topics)
run through notify.evaluate with an in-memory state; every delivered
notification is recorded. The Rust test replays the same steps (RID readings
only, so the ThaiWater additions stay out of play) and must deliver exactly
the same messages.
"""

import datetime
import gzip
import json
import random
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parents[2]))

from src import notify  # noqa: E402


class Publisher(notify.NtfyPublisher):
    def __init__(self):
        super().__init__("http://ntfy.invalid", prefix="ping")
        self.down = set()

    def publish(self, n):
        return n.topic not in self.down


STATIONS = {  # around (warning, danger)
    "P.1": (3.7, 4.2),
    "P.77": (4.3, 4.9),
    "P.21": (3.2, 3.6),
    "P.20": (2.35, 2.8),
}


def main():
    rng = random.Random(20261018)
    sequences = []
    for _ in range(40):
        state = notify.InMemoryState()
        pub = Publisher()
        t = datetime.datetime(2026, 9, 1) + datetime.timedelta(
            hours=rng.randrange(2000)
        )
        codes = rng.sample(sorted(STATIONS), rng.randint(1, 4))
        levels = {c: STATIONS[c][0] + rng.uniform(-0.6, 0.6) for c in codes}
        p_out = rng.random()
        steps = []
        for _ in range(rng.randint(8, 25)):
            gap = rng.choice([1, 1, 1, 1, 2, 4, 6])
            t += datetime.timedelta(hours=gap)
            # the clock sometimes runs ahead of the data (a stale feed)
            now = t + datetime.timedelta(minutes=rng.choice([5, 5, 30, 200, 300]))
            readings = []
            for c in codes:
                levels[c] += rng.uniform(-0.25, 0.25)
                pct = rng.choice([None, rng.uniform(5, 120)])
                level = round(levels[c], 2) if rng.random() > 0.05 else None
                readings.append(
                    {
                        "station_code": c,
                        "timestamp": t,
                        "water_level": level,
                        "discharge_percent": None if pct is None else round(pct, 2),
                    }
                )
            p_out = min(1.0, max(0.0, p_out + rng.uniform(-0.3, 0.3)))
            forecasts = [
                {
                    "station_code": "P.1",
                    "horizon_hours": rng.choice([24, 24, 24, 12]),
                    "p_warning": p_out,
                    "predicted_max_level": round(rng.uniform(2.5, 5.0), 3),
                    "source": rng.choice(["model", "model", "model", "heuristic"]),
                }
            ]
            pub.down = set(
                rng.sample(
                    [
                        "ping-p1-warning",
                        "ping-warning",
                        "ping-status",
                        "ping-p1-outlook",
                    ],
                    rng.choice([0, 0, 0, 1]),
                )
            )
            sent = notify.evaluate(readings, forecasts, state, pub, now=now)
            steps.append(
                {
                    "now": now.isoformat(),
                    "readings": [
                        {**r, "timestamp": r["timestamp"].isoformat()} for r in readings
                    ],
                    "forecasts": forecasts,
                    "down": sorted(pub.down),
                    "sent": [
                        {
                            "topic": n.topic,
                            "title": n.title,
                            "message": n.message,
                            "priority": n.priority,
                            "tags": n.tags or [],
                        }
                        for n in sent
                    ],
                }
            )
        sequences.append(steps)
    out = HERE / "notify_python_cases.json.gz"
    with open(out, "wb") as raw, gzip.GzipFile(fileobj=raw, mode="wb", mtime=0) as gz:
        gz.write(json.dumps({"sequences": sequences}).encode())
    total = sum(len(s["sent"]) for seq in sequences for s in seq)
    print(
        f"{len(sequences)} sequences, {sum(map(len, sequences))} steps, {total} notifications"
    )


if __name__ == "__main__":
    main()