test(QA): refactor Locust tests, add rate-limit and S3 load tests

Makefile

@@ -1,6 +1,6 @@
 HOST ?= http://localhost:8080
 
-.PHONY: stress-test oversell-test orders-test
+.PHONY: stress-test oversell-test orders-test mixed-test ratelimit-test s3-test
 
 stress-test:
 	uv run locust \
@@ -29,3 +29,30 @@ orders-test:
 		--spawn-rate 50 \
 		--run-time 60s \
 		--host $(HOST)
+
+mixed-test:
+	uv run locust \
+		-f load_tests/locustfile_mixed.py \
+		--headless \
+		--users 100 \
+		--spawn-rate 50 \
+		--run-time 60s \
+		--host $(HOST)
+
+ratelimit-test:
+	uv run locust \
+		-f load_tests/locustfile_ratelimit.py \
+		--headless \
+		--users 1 \
+		--spawn-rate 1 \
+		--run-time 15s \
+		--host $(HOST)
+
+s3-test:
+	uv run locust \
+		-f load_tests/locustfile_s3.py \
+		--headless \
+		--users 50 \
+		--spawn-rate 10 \
+		--run-time 20s \
+		--host $(HOST)

load_tests/locust_base.py

@@ -15,6 +15,7 @@ class BaseUser(HttpUser):
 
     def on_start(self) -> None:
         self.access_token: str | None = None
+        self.refresh_token: str | None = None
         email = f'locust_{uuid4()}@mail.com'
         password = '12345'
 
@@ -37,7 +38,9 @@ def on_start(self) -> None:
             catch_response=True,
         ) as token_res:
             if token_res.status_code == HTTPStatus.OK:
-                self.access_token = token_res.json().get('access_token')
+                token_data = token_res.json()
+                self.access_token = token_data.get('access_token')
+                self.refresh_token = token_data.get('refresh_token')
             else:
                 token_res.failure(f'Login failed: {token_res.status_code}')
 
@@ -47,3 +50,116 @@ def auth_headers(self) -> dict[str, str]:
             'Authorization': f'Bearer {self.access_token}',
             'X-Idempotency-Key': str(uuid4()),
         }
+
+    def do_get_presigned_url(
+        self,
+        product_id: str,
+        filename: str,
+        content_type: str,
+    ) -> str | None:
+        with self.client.post(
+            f'/api/v1/media/products/{product_id}/upload_url',
+            headers=self.auth_headers,
+            json={'filename': filename, 'content_type': content_type},
+            catch_response=True,
+        ) as res:
+            if res.status_code in (
+                HTTPStatus.OK,
+                HTTPStatus.TOO_MANY_REQUESTS,
+                HTTPStatus.BAD_REQUEST,
+            ):
+                res.success()
+                try:
+                    url = res.json().get('presigned_url')
+                    return str(url) if url else None
+                except Exception:
+                    return None
+            else:
+                res.failure(f'Get presigned url failed: {res.status_code}')
+                return None
+
+    def do_reserve(self, product_id: str, quantity: int = 1) -> str | None:
+        with self.client.post(
+            '/api/v1/inventory/reserve',
+            headers=self.auth_headers,
+            json={'product_id': product_id, 'quantity': quantity},
+            catch_response=True,
+        ) as res:
+            if res.status_code in (
+                HTTPStatus.OK,
+                HTTPStatus.CREATED,
+                HTTPStatus.CONFLICT,
+                HTTPStatus.TOO_MANY_REQUESTS,
+                HTTPStatus.BAD_REQUEST,
+            ):
+                res.success()
+                if res.status_code in (HTTPStatus.OK, HTTPStatus.CREATED):
+                    try:
+                        res_id = res.json().get('id')
+                        return str(res_id) if res_id is not None else None
+                    except Exception:
+                        pass
+                return None
+            else:
+                res.failure(f'Reserve failed: {res.status_code}')
+                return None
+
+    def do_create_order(self, reservation_id: str) -> bool:
+        with self.client.post(
+            '/api/v1/orders/',
+            headers=self.auth_headers,
+            json={'reservation_id': reservation_id},
+            catch_response=True,
+        ) as res:
+            if res.status_code in (
+                HTTPStatus.CREATED,
+                HTTPStatus.OK,
+                HTTPStatus.TOO_MANY_REQUESTS,
+            ):
+                res.success()
+                return True
+            else:
+                res.failure(f'Order failed: {res.status_code}')
+                return False
+
+    def do_view_profile(self) -> bool:
+        with self.client.get(
+            '/api/v1/users/me',
+            headers=self.auth_headers,
+            catch_response=True,
+        ) as res:
+            if res.status_code in (
+                HTTPStatus.OK,
+                HTTPStatus.TOO_MANY_REQUESTS,
+                HTTPStatus.BAD_REQUEST,
+            ):
+                res.success()
+                return True
+            else:
+                res.failure(f'Profile view failed: {res.status_code}')
+                return False
+
+    def do_refresh_token(self) -> bool:
+        if not self.refresh_token:
+            return False
+        with self.client.post(
+            '/api/v1/auth/refresh',
+            headers=self.auth_headers,
+            json={'refresh_token': self.refresh_token},
+            catch_response=True,
+        ) as res:
+            if res.status_code == HTTPStatus.OK:
+                token_data = res.json()
+                self.access_token = token_data.get('access_token')
+                self.refresh_token = token_data.get('refresh_token')
+                res.success()
+                return True
+            elif res.status_code in (
+                HTTPStatus.TOO_MANY_REQUESTS,
+                HTTPStatus.BAD_REQUEST,
+            ):
+                res.success()
+                return False
+            else:
+                res.failure(f'Token refresh failed: {res.status_code}')
+                return False

load_tests/locustfile.py

@@ -1,5 +1,3 @@
-from http import HTTPStatus
-
 from locust import between, task
 
 from load_tests.locust_base import BaseUser
@@ -14,17 +12,4 @@ class HighLoadUser(BaseUser):
     def reserve_product(self) -> None:
         if not self.access_token:
             return
-        with self.client.post(
-            '/api/v1/inventory/reserve',
-            headers=self.auth_headers,
-            json={'product_id': TARGET_PRODUCT_ID, 'quantity': 1},
-            catch_response=True,
-        ) as reserve_res:
-            if reserve_res.status_code not in (
-                HTTPStatus.OK,
-                HTTPStatus.CREATED,
-                HTTPStatus.CONFLICT,
-                HTTPStatus.TOO_MANY_REQUESTS,
-                HTTPStatus.BAD_REQUEST,
-            ):
-                reserve_res.failure(f'Reserve failed: {reserve_res.status_code}')
+        self.do_reserve(TARGET_PRODUCT_ID)

load_tests/locustfile_mixed.py

@@ -0,0 +1,27 @@
+from locust import between, task
+
+from load_tests.locust_base import BaseUser
+
+TARGET_PRODUCT_ID = '5995fa75-07c7-4b55-82b7-6bfbb52948b8'
+
+
+class MixedLoadUser(BaseUser):
+    wait_time = between(1.0, 3.0)
+
+    @task(6)
+    def reserve_product(self) -> None:
+        if not self.access_token:
+            return
+        self.do_reserve(TARGET_PRODUCT_ID)
+
+    @task(2)
+    def view_profile(self) -> None:
+        if not self.access_token:
+            return
+        self.do_view_profile()
+
+    @task(2)
+    def update_token(self) -> None:
+        if not self.refresh_token:
+            return
+        self.do_refresh_token()

load_tests/locustfile_orders.py

@@ -1,5 +1,3 @@
-from http import HTTPStatus
-
 from locust import between, task
 
 from load_tests.locust_base import BaseUser
@@ -14,29 +12,6 @@ class OrderLoadUser(BaseUser):
     def reserve_and_order(self) -> None:
         if not self.access_token:
             return
-        with self.client.post(
-            '/api/v1/inventory/reserve',
-            headers=self.auth_headers,
-            json={'product_id': TARGET_PRODUCT_ID, 'quantity': 1},
-            catch_response=True,
-        ) as response:
-            if response.status_code not in (
-                HTTPStatus.OK,
-                HTTPStatus.CREATED,
-                HTTPStatus.CONFLICT,
-                HTTPStatus.TOO_MANY_REQUESTS,
-                HTTPStatus.BAD_REQUEST,
-            ):
-                response.failure(f'Reserve failed: {response.status_code}')
-                return
-            if response.status_code not in (HTTPStatus.OK, HTTPStatus.CREATED):
-                return
-            reservation_id = response.json()['id']
-        with self.client.post(
-            '/api/v1/orders/',
-            headers=self.auth_headers,
-            json={'reservation_id': reservation_id},
-            catch_response=True,
-        ) as response:
-            if response.status_code not in (HTTPStatus.CREATED, HTTPStatus.OK):
-                response.failure(f'Order failed: {response.status_code}')
+        reservation_id = self.do_reserve(TARGET_PRODUCT_ID)
+        if reservation_id:
+            self.do_create_order(reservation_id)

load_tests/locustfile_oversell.py

@@ -1,5 +1,3 @@
-from http import HTTPStatus
-
 from locust import between, task
 
 from load_tests.locust_base import BaseUser
@@ -14,17 +12,4 @@ class OversellTestUser(BaseUser):
     def reserve_oversell_product(self) -> None:
         if not self.access_token:
             return
-        with self.client.post(
-            '/api/v1/inventory/reserve',
-            headers=self.auth_headers,
-            json={'product_id': OVERSELL_PRODUCT_ID, 'quantity': 1},
-            catch_response=True,
-        ) as reserve_res:
-            if reserve_res.status_code not in (
-                HTTPStatus.OK,
-                HTTPStatus.CREATED,
-                HTTPStatus.CONFLICT,
-                HTTPStatus.TOO_MANY_REQUESTS,
-                HTTPStatus.BAD_REQUEST,
-            ):
-                reserve_res.failure(f'Reserve failed: {reserve_res.status_code}')
+        self.do_reserve(OVERSELL_PRODUCT_ID)

load_tests/locustfile_ratelimit.py

@@ -0,0 +1,15 @@
+from locust import constant_pacing, task
+
+from load_tests.locust_base import BaseUser
+
+TARGET_PRODUCT_ID = '5995fa75-07c7-4b55-82b7-6bfbb52948b8'
+
+
+class RatelimitUser(BaseUser):
+    wait_time = constant_pacing(0.05)
+
+    @task
+    def spam_profile(self) -> None:
+        if not self.access_token:
+            return
+        self.do_reserve(TARGET_PRODUCT_ID)

load_tests/locustfile_s3.py

@@ -0,0 +1,13 @@
+from locust import between, task
+
+from load_tests.locust_base import BaseUser
+
+TARGET_PRODUCT_ID = '5995fa75-07c7-4b55-82b7-6bfbb52948b8'
+
+
+class PreSignedUrlForS3User(BaseUser):
+    wait_time = between(1.0, 3.0)
+
+    @task
+    def get_presigned_url(self) -> None:
+        self.do_get_presigned_url(TARGET_PRODUCT_ID, 'test.jpg', 'image/jpeg')