docqube_backend/scripts/loadtest.py

"""
B6.1 — drive concurrent load against a running DocQube and report percentiles.

Deliberately a script, not a test. Latency depends on the machine, the disk and
whatever else is running, so asserting a threshold in CI produces a suite that
fails at random — and a suite that fails at random gets switched off. The
per-request *query counts* are the CI gate (`test_query_budget.py`); this is for
answering "what happens at 50 concurrent users?" before a customer does.

    # against a locally running app
    APP_ENV=rework python scripts/loadtest.py --users 20 --requests 10

    # a single endpoint, more load
    APP_ENV=rework python scripts/loadtest.py --users 50 --requests 20 \\
        --only /api/drive/root

It mints its own token from `APP_SECRET`, so it needs no password and no seeded
user beyond one that exists. **Point it at staging, never production**: it
issues real requests and the write paths are excluded only because they are not
listed, not because anything stops them.
"""

import argparse
import os
import statistics
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

os.environ.setdefault("APP_ENV", "rework")
os.environ.setdefault("PYTHONIOENCODING", "utf-8")
sys.path.insert(0, str(Path(__file__).parent.parent))

# Read-only endpoints. Adding a mutation here means every run writes to the
# target — deliberately not done.
DEFAULT_PATHS = [
    "/api/health",
    "/api/auth/me",
    "/api/me/profile",
    "/api/notifications",
    "/api/notifications/unread-count",
    "/api/activity-logs",
    "/api/storage/usage",
    "/api/storage/files",
    "/api/drive/root",
    "/api/drive/me/recent",
]


def percentile(values: list[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    k = (len(ordered) - 1) * pct
    lo, hi = int(k), min(int(k) + 1, len(ordered) - 1)
    return ordered[lo] + (ordered[hi] - ordered[lo]) * (k - lo)


def mint_token(user_id: int, tenant_id: str | None) -> str:
    from app.core.security import create_access_token

    payload = {"sub": str(user_id)}
    if tenant_id:
        payload["tenant_id"] = tenant_id
    return create_access_token(payload)


def find_a_user() -> tuple[int, str | None]:
    """Any active user will do — this measures the endpoint, not the account."""
    from sqlalchemy import create_engine, text

    from app.core.settings import settings

    engine = create_engine(settings.DATABASE_URL)
    with engine.connect() as conn:
        row = conn.execute(
            text(
                "SELECT id, tenant_id FROM users "
                "WHERE is_deleted = false AND tenant_id IS NOT NULL "
                "ORDER BY id LIMIT 1"
            )
        ).first()
    engine.dispose()
    if row is None:
        raise SystemExit(
            "No tenant user found. Seed the database first: python scripts/seed.py"
        )
    return int(row[0]), str(row[1])


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default="http://localhost:20001")
    parser.add_argument("--users", type=int, default=10, help="concurrent workers")
    parser.add_argument("--requests", type=int, default=10, help="requests per worker")
    parser.add_argument("--only", action="append", help="restrict to these paths")
    args = parser.parse_args()

    import httpx

    paths = args.only or DEFAULT_PATHS
    user_id, tenant_id = find_a_user()
    headers = {"Authorization": f"Bearer {mint_token(user_id, tenant_id)}"}

    print(f"target      {args.base_url}")
    print(f"concurrency {args.users} workers x {args.requests} requests")
    print(f"endpoints   {len(paths)}")
    print()

    results: dict[str, list[float]] = {p: [] for p in paths}
    errors: dict[str, int] = {p: 0 for p in paths}

    def worker(_: int) -> None:
        with httpx.Client(base_url=args.base_url, timeout=30.0) as client:
            for _ in range(args.requests):
                for path in paths:
                    started = time.perf_counter()
                    try:
                        response = client.get(path, headers=headers)
                        elapsed = (time.perf_counter() - started) * 1000
                        if response.status_code >= 500:
                            errors[path] += 1
                        results[path].append(elapsed)
                    except Exception:  # noqa: BLE001 - a failure is a data point
                        errors[path] += 1

    started = time.perf_counter()
    with ThreadPoolExecutor(max_workers=args.users) as pool:
        list(pool.map(worker, range(args.users)))
    wall = time.perf_counter() - started

    total = sum(len(v) for v in results.values())
    print(f"{'endpoint':<38}{'n':>6}{'p50':>9}{'p95':>9}{'p99':>9}{'max':>9}{'5xx':>6}")
    print("-" * 86)
    for path in paths:
        samples = results[path]
        if not samples:
            print(f"{path:<38}{'—':>6}")
            continue
        print(
            f"{path:<38}{len(samples):>6}"
            f"{statistics.median(samples):>8.0f}ms"
            f"{percentile(samples, 0.95):>8.0f}ms"
            f"{percentile(samples, 0.99):>8.0f}ms"
            f"{max(samples):>8.0f}ms"
            f"{errors[path]:>6}"
        )

    print("-" * 86)
    print(f"{total} requests in {wall:.1f}s — {total / wall:.0f} req/s")

    if any(errors.values()):
        print("\nServer errors occurred; the percentiles above exclude them.")
        return 1
    return 0


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