From 090671b5ba7e20184c01d66247860487be0c593a Mon Sep 17 00:00:00 2001
From: David MENDY <davebull.dm@gmail.com>
Date: Tue, 7 Jan 2025 15:16:15 +0100
Subject: [PATCH] test: :white_check_mark: add api rates limit tests

---
 .github/workflows/main-ci.yml |  1 +
 backend/requirements.txt      |  3 ++-
 backend/src/router.py         |  6 +-----
 backend/tests/locust.conf     |  6 ++++++
 backend/tests/locustfile.py   | 18 ++++++++++++++++++
 docker-compose.ci.yml         |  2 --
 6 files changed, 28 insertions(+), 8 deletions(-)
 create mode 100644 backend/tests/locust.conf
 create mode 100644 backend/tests/locustfile.py

diff --git a/.github/workflows/main-ci.yml b/.github/workflows/main-ci.yml
index 53181227..a8fe191b 100644
--- a/.github/workflows/main-ci.yml
+++ b/.github/workflows/main-ci.yml
@@ -54,6 +54,7 @@ jobs:
         run: IMAGE_TAG=${{ needs.build-backend.outputs.image_tag }} docker compose -f docker-compose.yml -f docker-compose.ci.yml up -d
       - run: docker compose exec backend coverage run -m pytest
       - run: docker compose exec backend coverage xml --ignore-errors
+      - run: docker compose exec backend locust -f tests/locustfile.py -t 5s --headless -H http://localhost:5000/api
       - name: Produce the coverage report
         uses: insightsengineering/coverage-action@v2
         with:
diff --git a/backend/requirements.txt b/backend/requirements.txt
index 0c4b85e7..f3aeea3d 100644
--- a/backend/requirements.txt
+++ b/backend/requirements.txt
@@ -18,4 +18,5 @@ numpy<3.0.0
 # Dev
 pytest==8.3.4
 coverage==7.6.10
-Faker==33.3.0
\ No newline at end of file
+Faker==33.3.0
+locust==2.32.5
\ No newline at end of file
diff --git a/backend/src/router.py b/backend/src/router.py
index f6b5ebd4..234c8a57 100644
--- a/backend/src/router.py
+++ b/backend/src/router.py
@@ -36,12 +36,8 @@
 )
 from .utils import get_current_user, send_mail, upload_image
 
-DISABLE_RATE_LIMITS = os.environ.get("DISABLE_RATE_LIMITS", "false").lower() == "true"
-
 router = APIRouter(prefix="/api")
-limiter = Limiter(
-    key_func=get_remote_address if not DISABLE_RATE_LIMITS else lambda: None
-)
+limiter = Limiter(key_func=get_remote_address)
 
 
 @router.get("/", response_class=PlainTextResponse)
diff --git a/backend/tests/locust.conf b/backend/tests/locust.conf
new file mode 100644
index 00000000..0c5bda81
--- /dev/null
+++ b/backend/tests/locust.conf
@@ -0,0 +1,6 @@
+locustfile = tests/locustfile.py
+headless = true
+host = http://localhost:5000/api
+run-time = 5s
+spawn-rate =
+users = 60
diff --git a/backend/tests/locustfile.py b/backend/tests/locustfile.py
new file mode 100644
index 00000000..4c58cfba
--- /dev/null
+++ b/backend/tests/locustfile.py
@@ -0,0 +1,18 @@
+import time
+
+from locust import HttpUser, task
+
+
+class TestLocust(HttpUser):
+    @task
+    def rate_limit_version(self):
+        self.client.get("/version")
+
+    @task
+    def test_upload(self):
+        with open("./tests/images/revolver.jpg", "rb") as f:
+            self.client.post(
+                "/upload",
+                files={"image": f},
+                data={"date": time.time()},
+            )
diff --git a/docker-compose.ci.yml b/docker-compose.ci.yml
index 179d17f3..c06d8291 100644
--- a/docker-compose.ci.yml
+++ b/docker-compose.ci.yml
@@ -1,5 +1,3 @@
 services:
   backend:
     image: ${IMAGE_TAG}
-    environment:
-      - DISABLE_RATE_LIMITS=true