benchmark/locustfile.py

import json
import os

from dotenv import load_dotenv
from locust import HttpUser, between, task

load_dotenv()


class ApiUser(HttpUser):
    host = 'http://127.0.0.1:8080/ml-api'
    wait_time = between(1, 2)

    @task
    def predict_all(self):
        headers = {
            'Api-key': os.getenv('API_KEY'),
            'Content-Type': 'application/json'
        }
        data = {
            "room_description": "room with sea view",
            "beds": [
                {"type": "double", "count": 2}
            ]
        }
        print("Headers:", headers)
        print("Data being sent:", data)
        with self.client.post("/predict/room/predict/all", headers=headers, data=json.dumps(data),  # noqa: E501
                              catch_response=True) as response:
            print("Response status code:", response.status_code)
            print("Response data:", response.text)
            if response.status_code != 200:
                response.failure(f"Failed with {response.status_code}: {response.text}")
            else:
                response.success()