tests/load/locustfile.py

# tests/load/locustfile.py
from locust import HttpUser, task, between
import os

API_KEY = os.getenv("API_KEY", "test-key")
TEST_REPO = os.getenv("LOAD_TEST_REPO", "https://github.com/psf/requests")


class CodeReviewUser(HttpUser):
    wait_time = between(2, 5)
    headers = {"X-API-Key": API_KEY, "Content-Type": "application/json"}

    @task(3)
    def submit_async_job(self):
        """High-frequency task: submit to async queue (cheap endpoint)."""
        with self.client.post(
            "/api/v1/analyze/async",
            json={"github_url": TEST_REPO},
            headers=self.headers,
            catch_response=True,
        ) as resp:
            if resp.status_code == 200:
                job_id = resp.json().get("job_id")
                if job_id:
                    resp.success()
                else:
                    resp.failure("No job_id in response")
            elif resp.status_code == 429:
                resp.success()  # Rate limit is expected under load
            else:
                resp.failure(f"Unexpected status {resp.status_code}")

    @task(1)
    def poll_health(self):
        """Lightweight health check to measure baseline latency."""
        self.client.get("/api/v1/health", headers=self.headers)

    @task(1)
    def poll_metrics(self):
        """Check internal metrics endpoint under load."""
        self.client.get("/metrics/prometheus")


class BurstUser(HttpUser):
    """Simulates burst traffic hitting the sync endpoint directly."""
    wait_time = between(5, 10)
    headers = {"X-API-Key": API_KEY, "Content-Type": "application/json"}

    @task(1)
    def sync_analyze(self):
        """Low-frequency task: full sync analyze (expensive, rate-limited)."""
        with self.client.post(
            "/api/v1/analyze",
            json={"github_url": TEST_REPO},
            headers=self.headers,
            catch_response=True,
            timeout=180,
        ) as resp:
            if resp.status_code in (200, 429):
                resp.success()
            else:
                resp.failure(f"Unexpected status {resp.status_code}")