Add hacky fix for rate-limiter segregation between pytests

This commit is contained in:
2024-11-09 23:23:55 -06:00
parent cd900288ff
commit 4c2c4bc2ad

View File

@@ -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(