| 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, |
| ) |
|
|
