File size: 1,575 Bytes
457c02a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
from typing import Optional
import requests
from locust import HttpUser, between, task
API_BASE_URL = "http://localhost:8000"
def login(username: str, password: str) -> Optional[str]:
"""This function calls the login endpoint of the API to authenticate the user and get a token.
Args:
username (str): email of the user
password (str): password of the user
Returns:
Optional[str]: token if login is successful, None otherwise
"""
url = f"{API_BASE_URL}/login"
headers = {
"accept": "application/json",
"Content-Type": "application/x-www-form-urlencoded",
}
data = {
"grant_type": "",
"username": username,
"password": password,
"scope": "",
"client_id": "",
"client_secret": "",
}
response = requests.post(url, headers=headers, data=data)
if response.status_code == 200:
return response.json()["access_token"]
else:
return None
class APIUser(HttpUser):
"""Represents a user in the API."""
wait_time = between(1, 5)
@task(1)
def predict(self):
"""Predicts an image using the model."""
token = login("admin@example.com", "admin")
files = [
("file", ("dog.jpeg", open("stress_test/dog.jpeg", "rb"), "image/jpeg"))
]
headers = {"Authorization": f"Bearer {token}"}
payload = {}
self.client.post(
"http://0.0.0.0:8000/model/predict",
headers=headers,
data=payload,
files=files,
)
|