cleaning and test fixes

Author: pedro-cf

stac_fastapi/core/stac_fastapi/core/rate_limit.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+from typing import Optional
 
 from fastapi import FastAPI, Request
 from slowapi import Limiter, _rate_limit_exceeded_handler
@@ -11,22 +12,32 @@
 
 logger = logging.getLogger(__name__)
 
-limiter = Limiter(key_func=get_remote_address)
+def get_limiter(key_func=get_remote_address):
+    return Limiter(key_func=key_func)
 
-
-def setup_rate_limit(app: FastAPI):
+def setup_rate_limit(
+    app: FastAPI, 
+    rate_limit: Optional[str] = None, 
+    key_func=get_remote_address
+):
     """Set up rate limiting middleware."""
-    RATE_LIMIT = os.getenv("STAC_FASTAPI_RATE_LIMIT")
+    RATE_LIMIT = rate_limit or os.getenv("STAC_FASTAPI_RATE_LIMIT")
+    
+    if not RATE_LIMIT:
+        logger.info("Rate limiting is disabled")
+        return
+
     logger.info(f"Setting up rate limit with RATE_LIMIT={RATE_LIMIT}")
-    if RATE_LIMIT:
-        app.state.limiter = limiter
-        app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
-        app.add_middleware(SlowAPIMiddleware)
-
-        @app.middleware("http")
-        @limiter.limit(RATE_LIMIT)
-        async def rate_limit_middleware(request: Request, call_next):
-            response = await call_next(request)
-            return response
+    
+    limiter = get_limiter(key_func)
+    app.state.limiter = limiter
+    app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
+    app.add_middleware(SlowAPIMiddleware)
+
+    @app.middleware("http")
+    @limiter.limit(RATE_LIMIT)
+    async def rate_limit_middleware(request: Request, call_next):
+        response = await call_next(request)
+        return response
 
     logger.info("Rate limit setup complete")

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/app.py

@@ -99,7 +99,7 @@
 app.root_path = os.getenv("STAC_FASTAPI_ROOT_PATH", "")
 
 # Add rate limit
-setup_rate_limit(app)
+setup_rate_limit(app, rate_limit=os.getenv("STAC_FASTAPI_RATE_LIMIT"))
 
 
 @app.on_event("startup")

stac_fastapi/opensearch/stac_fastapi/opensearch/app.py

@@ -99,7 +99,7 @@
 app.root_path = os.getenv("STAC_FASTAPI_ROOT_PATH", "")
 
 # Add rate limit
-setup_rate_limit(app)
+setup_rate_limit(app, rate_limit=os.getenv("STAC_FASTAPI_RATE_LIMIT"))
 
 
 @app.on_event("startup")

stac_fastapi/tests/conftest.py

@@ -292,9 +292,7 @@ async def app_rate_limit():
     ).app
 
     # Set up rate limit
-    os.environ["STAC_FASTAPI_RATE_LIMIT"] = "2/minute"
-    setup_rate_limit(app)
-    del os.environ["STAC_FASTAPI_RATE_LIMIT"]
+    setup_rate_limit(app, rate_limit="2/minute")
 
     return app