mirror of
https://github.com/Xevion/linkpulse.git
synced 2025-12-06 03:15:35 -06:00
Add hacky fix for rate-limiter segregation between pytests
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user