Northern-Thailand-Ping-River-Monitor/scripts/load_test.py

#!/usr/bin/env python3
"""Staged load / client-stress test for the Ping River Monitor API + dashboard.

Simulates a realistic traffic mix (dashboard page loads, the API calls the
dashboard itself makes, heavy history queries, external API consumers) at
increasing concurrency stages, and reports throughput, latency percentiles,
and errors per stage plus the slowest endpoints.

Run against a LOCAL instance for full stress (never full-stress production —
it hosts live flood monitoring):

    python -m uvicorn src.web_api:app --port 8125   # separate shell
    python scripts/load_test.py http://localhost:8125

A gentle production baseline (low, fixed concurrency):

    python scripts/load_test.py https://water.buildfor.life --gentle
"""

import argparse
import random
import statistics
import threading
import time
from collections import Counter

import requests

# Weighted endpoint mix: dashboard session + API consumers
ENDPOINTS = [
    ("/", 10),
    ("/measurements/latest?limit=500", 20),
    ("/stations", 10),
    ("/api/hii/rainfall/latest", 15),
    ("/api/hii/waterlevel/latest", 15),
    ("/forecast", 10),
    ("/api/stats", 5),
    ("/measurements/history/P.1?hours=168", 10),
    ("/measurements/history/P.67?hours=720", 5),
    ("/health", 5),
]
POOL = [endpoint for endpoint, weight in ENDPOINTS for _ in range(weight)]

FULL_STAGES = [(10, 20), (50, 20), (200, 25)]  # (clients, seconds)
GENTLE_STAGES = [(3, 15), (8, 15)]


def _worker(base, stop_at, results, errors):
    session = requests.Session()
    while time.time() < stop_at:
        path = random.choice(POOL)
        start = time.perf_counter()
        try:
            response = session.get(f"{base}{path}", timeout=30)
            elapsed = time.perf_counter() - start
            if response.status_code == 200:
                results.append((path, elapsed))
            else:
                errors.append((path, response.status_code))
        except Exception as error:
            errors.append((path, type(error).__name__))


def _pct(values, p):
    if len(values) >= 100:
        return statistics.quantiles(values, n=100)[p - 1]
    return max(values)


def run_stage(base, clients, seconds):
    results, errors = [], []
    stop_at = time.time() + seconds
    threads = [
        threading.Thread(
            target=_worker, args=(base, stop_at, results, errors), daemon=True
        )
        for _ in range(clients)
    ]
    for thread in threads:
        thread.start()
    for thread in threads:
        thread.join(timeout=seconds + 35)

    latencies = [elapsed for _, elapsed in results]
    total = len(results) + len(errors)
    print(f"\n== {clients} clients x {seconds}s ==")
    print(
        f"requests: {total}  ok: {len(results)}  errors: {len(errors)}  "
        f"rps: {total / seconds:.1f}"
    )
    if latencies:
        print(
            f"latency ms  p50: {statistics.median(latencies) * 1000:.0f}  "
            f"p95: {_pct(latencies, 95) * 1000:.0f}  "
            f"p99: {_pct(latencies, 99) * 1000:.0f}  "
            f"max: {max(latencies) * 1000:.0f}"
        )
    by_endpoint = {}
    for path, elapsed in results:
        by_endpoint.setdefault(path, []).append(elapsed)
    slowest = sorted(
        by_endpoint.items(), key=lambda kv: -statistics.median(kv[1])
    )[:4]
    for path, values in slowest:
        print(
            f"  slow: {path:45} n={len(values):5} "
            f"p50={statistics.median(values) * 1000:6.0f}ms "
            f"max={max(values) * 1000:7.0f}ms"
        )
    if errors:
        top = Counter(f"{path} {code}" for path, code in errors).most_common(5)
        print(f"  errors: {top}")
    return {"clients": clients, "total": total, "errors": len(errors)}


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("base", nargs="?", default="http://localhost:8125")
    parser.add_argument(
        "--gentle",
        action="store_true",
        help="low fixed concurrency (safe for the production instance)",
    )
    args = parser.parse_args(argv)
    base = args.base.rstrip("/")

    # Warm caches first so stage 1 doesn't measure cold-start work
    for path in ("/forecast", "/api/stats", "/measurements/latest?limit=500"):
        try:
            requests.get(f"{base}{path}", timeout=60)
        except Exception:
            pass

    print(f"target: {base}  mode: {'gentle' if args.gentle else 'full'}")
    stages = GENTLE_STAGES if args.gentle else FULL_STAGES
    summary = [run_stage(base, clients, seconds) for clients, seconds in stages]
    worst = max(
        (stage["errors"] / stage["total"] for stage in summary if stage["total"]),
        default=1.0,
    )
    print(f"\nworst-stage error rate: {worst:.1%}")
    return 0 if worst < 0.05 else 1


if __name__ == "__main__":
    import sys

    sys.exit(main())