From 4c2c4bc2adc8ae7dca6ee60c018cee5553b03018 Mon Sep 17 00:00:00 2001 From: Xevion Date: Sat, 9 Nov 2024 23:23:55 -0600 Subject: [PATCH] Add hacky fix for rate-limiter segregation between pytests --- backend/linkpulse/dependencies.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/backend/linkpulse/dependencies.py b/backend/linkpulse/dependencies.py index 514b48b..1df57c5 100644 --- a/backend/linkpulse/dependencies.py +++ b/backend/linkpulse/dependencies.py @@ -1,3 +1,4 @@ +import os import structlog from fastapi import HTTPException, Request, Response, status from limits.strategies import MovingWindowRateLimiter @@ -8,6 +9,7 @@ storage = MemoryStorage() strategy = MovingWindowRateLimiter(storage) logger = structlog.get_logger() +is_pytest = os.environ.get("PYTEST_VERSION") is not None class RateLimiter: @@ -16,15 +18,19 @@ class RateLimiter: self.retry_after = str(self.limit.get_expiry()) async def __call__(self, request: Request, response: Response): - logger.debug("Rate limiting request", path=request.url.path) - key = request.headers.get("X-Real-IP") + if key is None: if request.client is None: logger.warning("No client information available for request.") return False key = request.client.host + if is_pytest: + # This is somewhat hacky, I'm not sure if there's a way it can break during pytesting, but look here if odd rate limiting errors occur during tests + # The reason for this is so tests don't compete with each other for rate limiting + key += "." + os.environ["PYTEST_CURRENT_TEST"] + if not strategy.hit(self.limit, key): logger.warning("Rate limit exceeded", key=key) raise HTTPException(