Scaffolding for a simple stress test

stress_test.py

import json
import random
from locust import HttpUser, task, between, events
from locust.exception import LocustError
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("--api-url", type=str, default="https://parallelum.com.br/fipe/api/v1/carros/marcas",
                    help="Full API endpoint URL to test (e.g., for brands)")
parser.add_argument("--request-body", type=str, default="{}",
                    help="JSON body as string (ignored for GET; for future POST extensions)")
args, _ = parser.parse_known_args()  # Ignore unknown args passed by Locust

try:
    if "/fipe/api/v1/carros/marcas" not in args.api_url:
        raise LocustError("Provided --api-url must include '/fipe/api/v1/carros/marcas'")
    base_host = args.api_url.split("/fipe/api/v1")[0]
    initial_path = "/fipe/api/v1/carros/marcas"
except Exception as e:
    raise LocustError(f"Invalid API URL: {e}")

try:
    request_body = json.loads(args.request_body)
except json.JSONDecodeError:
    request_body = {}
    events.request_failure.fire(request_type="INIT", name="Parse Body", response_time=0, exception="Invalid JSON body", response_length=0)

class FIPEAPIUser(HttpUser):
    host = base_host
    wait_time = between(1, 5)  # Randomized wait

    # Shared state across tasks
    brands = []
    models = {}  # Dict: brand_codigo -> list of models
    years = {}   # Dict: (brand_codigo, model_codigo) -> list of years

    def on_start(self):
        """Fetch and cache brands on user start for chained requests."""
        response = self.client.get(initial_path, name="Get Brands (Init)")
        if response.status_code == 200:
            try:
                self.brands = response.json()
                events.request_success.fire(request_type="GET", name="Get Brands (Init)", response_time=response.elapsed.total_seconds() * 1000, response_length=len(response.content))
            except json.JSONDecodeError:
                events.request_failure.fire(request_type="GET", name="Get Brands (Init)", response_time=response.elapsed.total_seconds() * 1000, exception="JSON Decode Error", response_length=0)
        else:
            events.request_failure.fire(request_type="GET", name="Get Brands (Init)", response_time=response.elapsed.total_seconds() * 1000, exception=f"Status {response.status_code}", response_length=0)

    @task(2)  # Higher weight: Basic endpoint, simulate frequent hits
    def get_brands(self):
        """Task: Fetch list of car brands."""
        self.client.get("/fipe/api/v1/carros/marcas", name="Get Brands")

    @task(3)  # Medium weight: Deeper navigation
    def get_models(self):
        """Task: Randomly select a brand and fetch its models."""
        if not self.brands:
            return  # Skip if init failed
        brand = random.choice(self.brands)
        codigo = brand["codigo"]
        response = self.client.get(f"/fipe/api/v1/carros/marcas/{codigo}/modelos", name="Get Models")
        if response.status_code == 200:
            try:
                modelos_data = response.json()["modelos"]
                self.models[codigo] = modelos_data  # Cache for further chaining
            except (json.JSONDecodeError, KeyError):
                pass  # Handled by Locust failure event

    @task(4)  # Higher weight: Even deeper, more realistic load on nested endpoints
    def get_years(self):
        """Task: Randomly select a brand and model, then fetch years."""
        if not self.models:
            return
        brand_codigo = random.choice(list(self.models.keys()))
        model = random.choice(self.models[brand_codigo])
        model_codigo = model["codigo"]
        self.client.get(f"/fipe/api/v1/carros/marcas/{brand_codigo}/modelos/{model_codigo}/anos", name="Get Years")

    @task(5)  # Highest weight: Full chain to value endpoint, simulates complete user flow
    def get_value(self):
        """Task: Randomly select brand, model, year, and fetch vehicle value."""
        if not self.brands or not self.models:
            return
        brand = random.choice(self.brands)
        brand_codigo = brand["codigo"]
        if brand_codigo not in self.models:
            return
        model = random.choice(self.models[brand_codigo])
        model_codigo = model["codigo"]
        # To avoid extra calls, simulate a year (API years are like '2023-1'; use a placeholder if not cached)
        # For realism, fetch years if not cached
        key = (brand_codigo, model_codigo)
        if key not in self.years:
            response = self.client.get(f"/fipe/api/v1/carros/marcas/{brand_codigo}/modelos/{model_codigo}/anos", name="Get Years (Sub)")
            if response.status_code == 200:
                try:
                    self.years[key] = response.json()
                except json.JSONDecodeError:
                    return
        if key in self.years and self.years[key]:
            year = random.choice(self.years[key])
            year_codigo = year["codigo"]
            self.client.get(f"/fipe/api/v1/carros/marcas/{brand_codigo}/modelos/{model_codigo}/anos/{year_codigo}", name="Get Value")

# Custom event hook for enhanced logging (optional, for customization)
@events.test_start.add_listener
def on_test_start(environment, **kwargs):
    print(f"Starting stress test against API URL: {args.api_url} with body: {args.request_body}")

@events.request.add_listener
def on_request(request_type, name, response_time, response_length, response, context, exception, start_time, url, **kwargs):
    if exception:
        print(f"Request failed: {name} - {exception}")