monitoring/locust/locustfile.py

"""
Locust Load Testing Script for Skill Classification API

This script defines user behavior for load testing the prediction and monitoring
endpoints of the Skill Classification API.
"""

from locust import HttpUser, task, between


class SkillClassificationUser(HttpUser):
    """
    Simulated user for load testing the Skill Classification API.
    
    This user performs the following actions:
    - Single predictions (most frequent)
    - Batch predictions
    - Monitoring and health checks
    """
    
    # Default host for the API (can be overridden via --host flag or Web UI)
    # HuggingFace Spaces URL (or use http://localhost:8080 for Docker, http://localhost:8000 for local dev)
    host = "https://dacrow13-hopcroft-skill-classification.hf.space"
    
    # Wait between 1 and 5 seconds between tasks to simulate real user behavior
    wait_time = between(1, 5)
    
    @task(3)
    def predict_single(self):
        """
        Task 1: Single Prediction (Weight: 3)
        
        Performs a POST request to /predict with a single issue text.
        This is the main task and executes more frequently due to higher weight.
        """
        payload = {
            "issue_text": "Fix authentication bug in login module"
        }
        
        with self.client.post(
            "/predict",
            json=payload,
            catch_response=True
        ) as response:
            if response.status_code == 201:
                response.success()
            else:
                response.failure(f"Prediction failed with status {response.status_code}")
    
    @task(1)
    def predict_batch(self):
        """
        Task 2: Batch Prediction (Weight: 1)
        
        Performs a POST request to /predict/batch with multiple issue texts.
        """
        payload = {
            "issues": [
                {"issue_text": "Test 1"},
                {"issue_text": "Test 2"}
            ]
        }
        
        with self.client.post(
            "/predict/batch",
            json=payload,
            catch_response=True
        ) as response:
            if response.status_code == 200:
                response.success()
            else:
                response.failure(f"Batch prediction failed with status {response.status_code}")
    
    @task(1)
    def monitoring_and_history(self):
        """
        Task 3: Monitoring and History (Weight: 1)
        
        Performs GET requests to check prediction history and system health.
        """
        # Check prediction history
        with self.client.get(
            "/predictions",
            catch_response=True
        ) as response:
            if 200 <= response.status_code < 300:
                response.success()
            else:
                response.failure(f"Predictions history failed with status {response.status_code}")
        
        # Check system health
        with self.client.get(
            "/health",
            catch_response=True
        ) as response:
            if 200 <= response.status_code < 300:
                response.success()
            else:
                response.failure(f"Health check failed with status {response.status_code}")