locustfile.py

"""
RoadSense AI — Locust Load Test Configuration
Tests throughput target: 1000 docs/hr (~17 signals/min)

Endpoints tested:
- POST /ingest-signal (primary load — signal ingestion)
- GET /incidents (read load — dashboard queries)
- GET /confidence/{id} (spot checks)

Usage:
  locust -f locustfile.py --host=https://api.roadsense.dev
  # Then open http://localhost:8089 to configure and start

  Headless mode (CI):
  locust -f locustfile.py --host=https://api.roadsense.dev \
         --users 10 --spawn-rate 2 --run-time 5m --headless

Target: 1000 signals/hour = ~17/min = ~0.28/sec
"""

import os
import uuid
import random
from datetime import datetime, timezone
from locust import HttpUser, task, between, tag


# ── Config ────────────────────────────────────────────────────────────────────

API_HOST = os.getenv("ROADSENSE_API_URL", "https://api.roadsense.dev")

SOURCES = ["reddit", "news", "weather"]
LANGUAGES = ["en", "hi", "ta", "te"]
LOCATIONS = [
    {"latitude": 12.9750, "longitude": 77.6060, "address": "MG Road"},
    {"latitude": 12.9719, "longitude": 77.6412, "address": "Indiranagar"},
    {"latitude": 12.9177, "longitude": 77.6229, "address": "Silk Board"},
    {"latitude": 12.9698, "longitude": 77.7500, "address": "Whitefield"},
    {"latitude": 12.8399, "longitude": 77.6770, "address": "Electronic City"},
    {"latitude": 12.9352, "longitude": 77.6245, "address": "Koramangala"},
]

ENGLISH_TEXTS = [
    "Huge pothole near {location}",
    "Road sinking at {location}",
    "Water logging spotted at {location}",
    "Cracks spreading on the road at {location}",
    "Dangerous dip near {location}",
    "Severe road damage around {location}",
    "Vehicles swerving to avoid potholes near {location}",
    "Standing water hiding potholes at {location}",
    "Traffic slowing due to damaged road at {location}",
    "Fresh cracks noticed at {location}",
]

HINDI_TEXTS = [
    "सड़क पर बड़ा गड्ढा है {location} में",
    "भारी बारिश से सड़क खराब हुई {location}",
    "पानी भरा है सड़क पर {location}",
    "सड़क पर दरारें बढ़ रही हैं {location}",
]

TAMIL_TEXTS = [
    "சாலையில் பெரிய குழி உள்ளது {location}",
    "மழையால் சாலை சேதமடைந்தது {location}",
    "சாலை சேதம் கடுமை {location}",
]

TELUGU_TEXTS = [
    "రోడ్డుపై పెద్ద గుంత ఉంది {location}",
    "భారీ వర్షంతో రోడ్డు దెబ్బతింది {location}",
    "గుంతలు ఎక్కువగా ఉన్నాయి {location}",
]

LANG_TEXTS = {
    "en": ENGLISH_TEXTS,
    "hi": HINDI_TEXTS,
    "ta": TAMIL_TEXTS,
    "te": TELUGU_TEXTS,
}


# ── Helpers ───────────────────────────────────────────────────────────────────

def generate_signal():
    """Generate a realistic random signal payload."""
    location = random.choice(LOCATIONS)
    language = random.choice(LANGUAGES)
    templates = LANG_TEXTS[language]
    content = random.choice(templates).format(location=location["address"])

    return {
        "signal_id": str(uuid.uuid4()),
        "original_content": content,
        "translated_content": content if language == "en" else None,
        "detected_language": language,
        "source": random.choice(SOURCES),
        "timestamp": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "location": location,
    }


# ── Locust User ──────────────────────────────────────────────────────────────

class RoadSenseUser(HttpUser):
    """
    Simulates a RoadSense API user.

    Wait time: 3-6 seconds between requests.
    With 10 users: ~2-3 requests/sec = ~7200-10800/hr (exceeds 1000/hr target).
    With 3 users:  ~0.5-1 req/sec = ~1800-3600/hr (closer to target).
    """
    wait_time = between(3, 6)

    @tag("ingest")
    @task(10)  # 10x weight — primary load
    def ingest_signal(self):
        """POST /ingest-signal — primary throughput test."""
        signal = generate_signal()
        with self.client.post(
            "/ingest-signal",
            json=signal,
            name="/ingest-signal",
            catch_response=True,
        ) as response:
            if response.status_code == 201:
                response.success()
            elif response.status_code == 400:
                response.failure(f"Bad request: {response.text[:100]}")
            else:
                response.failure(f"Unexpected status: {response.status_code}")

    @tag("read")
    @task(3)  # 3x weight — dashboard reads
    def get_incidents(self):
        """GET /incidents — simulates dashboard polling."""
        with self.client.get(
            "/incidents",
            name="/incidents",
            catch_response=True,
        ) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"Status: {response.status_code}")

    @tag("read")
    @task(1)  # 1x weight — occasional confidence check
    def get_confidence(self):
        """GET /confidence/{id} — simulates spot-check."""
        # Use a fake ID — expect 404, just testing the endpoint handles it
        fake_id = str(uuid.uuid4())
        with self.client.get(
            f"/confidence/{fake_id}",
            name="/confidence/[id]",
            catch_response=True,
        ) as response:
            if response.status_code in [200, 404]:
                response.success()
            else:
                response.failure(f"Status: {response.status_code}")

    @tag("export")
    @task(1)  # 1x weight — rare export
    def export_data(self):
        """GET /export — simulates data export."""
        with self.client.get(
            "/export",
            name="/export",
            catch_response=True,
        ) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"Status: {response.status_code}")