scripts/locustfile.py

"""Locust Load Testing Suite for Real-Time Inference Gateway & API Layer.

Simulates concurrent banking client nodes, automated payment streams,
and compliance investigators to benchmark SLA adherence (<100ms p99 latency),
throughput capacity (TPS), and error rates.

Usage:
    # Headless 60-second load test with 50 concurrent users
    locust -f scripts/locustfile.py --headless -u 50 -r 10 --run-time 60s --host http://localhost:8000

    # Web UI mode
    locust -f scripts/locustfile.py --host http://localhost:8000
"""

from __future__ import annotations

# gevent monkey-patching: ensure socket and ssl are safely patched before standard networking modules
# to prevent RecursionError on Python 3.12 under high-concurrency load tests.
try:
    import gevent.monkey
    if not gevent.monkey.is_module_patched("socket"):
        gevent.monkey.patch_all()
except (ImportError, Exception):
    pass

import random
import uuid
from typing import Any

try:
    from locust import FastHttpUser, between, task
except ImportError:
    class FastHttpUser:  # type: ignore[no-redef]
        pass

    def between(a: float, b: float) -> Any:  # type: ignore[no-redef]
        return None

    def task(weight: int = 1) -> Any:  # type: ignore[no-redef]
        def decorator(f: Any) -> Any:
            return f
        return decorator

_BANK_IDS = ["bank_alpha", "bank_beta", "bank_gamma"]
_CURRENCIES = ["EUR", "USD", "GBP", "CHF"]
_MERCHANTS = ["crypto_exchange", "electronics", "wire_transfer", "gambling", "retail", "jewelry"]


class BankingNodeUser(FastHttpUser):
    """Simulates active banking nodes streaming real-time transactions for scoring."""

    wait_time = between(0.01, 0.05)  # High frequency payment stream (20-100 req/s per user)

    @task(60)
    def score_realtime_transaction(self) -> None:
        """Benchmark /api/v1/predict endpoint for real-time <100ms inference SLA."""
        tx_id = f"txn_{uuid.uuid4().hex[:12]}"
        bank_id = random.choice(_BANK_IDS)
        amount = round(random.uniform(5.0, 50000.0), 2)
        merchant = random.choice(_MERCHANTS)

        payload: dict[str, Any] = {
            "transaction_id": tx_id,
            "bank_id": bank_id,
            "source_bank_id": bank_id,
            "amount": amount,
            "currency": random.choice(_CURRENCIES),
            "merchant_category": merchant,
            "sender_account": f"DE893704004405320130{random.randint(10, 99)}",
            "receiver_account": f"GB29NWBK601613319268{random.randint(10, 99)}",
            "device_fingerprint": f"dev_fp_{uuid.uuid4().hex[:8]}",
            "country_code": random.choice(["DE", "FR", "US", "GB", "NL", "TR"]),
            "timestamp": "2026-09-02T05:00:00Z",
        }

        headers = {
            "X-Tenant-ID": bank_id,
            "X-Bank-ID": bank_id,
            "Content-Type": "application/json",
        }

        with self.client.post(
            "/api/v1/predict",
            json=payload,
            headers=headers,
            name="POST /api/v1/predict",
            catch_response=True,
        ) as response:
            if response.status_code == 200:
                data = response.json()
                if "risk_score" in data or "composite_risk_score" in data:
                    response.success()
                else:
                    response.failure(f"Missing risk score in payload: {response.text}")
            elif response.status_code == 429:
                # Rate limited under extreme burst (expected at high concurrency)
                response.success()
            else:
                response.failure(f"Unexpected status code: {response.status_code}")

    @task(20)
    def score_v1_inference_endpoint(self) -> None:
        """Benchmark /v1/inference/score endpoint."""
        bank_id = random.choice(_BANK_IDS)
        payload = {
            "transaction_id": f"txn_{uuid.uuid4().hex[:12]}",
            "source_bank_id": bank_id,
            "amount": round(random.uniform(10.0, 150000.0), 2),
            "currency": "EUR",
            "sender_account": f"DE893704004405320130{random.randint(10, 99)}",
            "receiver_account": f"GB29NWBK601613319268{random.randint(10, 99)}",
            "merchant_category": random.choice(_MERCHANTS),
            "device_fingerprint": f"dev_fp_{uuid.uuid4().hex[:8]}",
        }

        self.client.post(
            "/v1/inference/score",
            json=payload,
            headers={"X-Tenant-ID": bank_id, "X-Bank-ID": bank_id},
            name="POST /v1/inference/score",
        )

    @task(10)
    def check_health_readiness(self) -> None:
        """Verify low-overhead liveness & readiness probes."""
        self.client.get("/health", name="GET /health")
        self.client.get("/ready", name="GET /ready")


class ComplianceAnalystUser(FastHttpUser):
    """Simulates compliance officers reviewing cases and authenticating."""

    wait_time = between(0.1, 0.5)

    def on_start(self) -> None:
        """Log in on session start."""
        self.auth_token = ""
        res = self.client.post(
            "/api/v1/auth/login",
            json={"username": "investigator_alpha", "password": "DemoPassword123!"},
            name="POST /api/v1/auth/login",
        )
        if res.status_code == 200:
            self.auth_token = res.json().get("access_token", "")

    @task(15)
    def fetch_cases_list(self) -> None:
        """Query case workbench cases with tenant isolation."""
        headers = {
            "Authorization": f"Bearer {self.auth_token}" if self.auth_token else "",
            "X-Tenant-ID": "bank_alpha",
            "X-Bank-ID": "bank_alpha",
        }
        self.client.get("/api/v1/cases", headers=headers, name="GET /api/v1/cases")

    @task(5)
    def check_lockout_status(self) -> None:
        """Query brute-force lockout status."""
        self.client.get(
            "/api/v1/auth/lockout-status?username=investigator_alpha",
            name="GET /api/v1/auth/lockout-status",
        )


if __name__ == "__main__":
    import sys
    try:
        from locust import main
        # Default to this locustfile if -f / --locustfile not explicitly given
        if "-f" not in sys.argv and "--locustfile" not in sys.argv:
            sys.argv.insert(1, "-f")
            sys.argv.insert(2, __file__)
        main.main()
    except ImportError:
        if "--help" in sys.argv or "-h" in sys.argv:
            print("locustfile.py: Load testing suite using Locust.")
            print("Usage: python scripts/locustfile.py [options]")
            print("\nNote: Locust is not installed in the current environment (pip install locust).")
            sys.exit(0)
        print("[ERROR] Locust is not installed. Run: pip install locust")
        sys.exit(1)