pre-commit

Author: pedro-cf

stac_fastapi/core/stac_fastapi/core/rate_limit.py

@@ -1,16 +1,21 @@
+"""Rate limiting middleware."""
+
+import logging
 import os
+
 from fastapi import FastAPI, Request
 from slowapi import Limiter, _rate_limit_exceeded_handler
-from slowapi.util import get_remote_address
 from slowapi.errors import RateLimitExceeded
 from slowapi.middleware import SlowAPIMiddleware
-import logging
+from slowapi.util import get_remote_address
 
 logger = logging.getLogger(__name__)
 
 limiter = Limiter(key_func=get_remote_address)
 
+
 def setup_rate_limit(app: FastAPI):
+    """Set up rate limiting middleware."""
     RATE_LIMIT = os.getenv("STAC_FASTAPI_RATE_LIMIT")
     logger.info(f"Setting up rate limit with RATE_LIMIT={RATE_LIMIT}")
     if RATE_LIMIT:
@@ -24,4 +29,4 @@ async def rate_limit_middleware(request: Request, call_next):
             response = await call_next(request)
             return response
 
-    logger.info("Rate limit setup complete")
+    logger.info("Rate limit setup complete")

stac_fastapi/elasticsearch/stac_fastapi/elasticsearch/app.py

@@ -1,7 +1,6 @@
 """FastAPI application."""
 
 import os
-from stac_fastapi.core.rate_limit import setup_rate_limit
 
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
@@ -18,6 +17,7 @@
     EsAsyncAggregationClient,
 )
 from stac_fastapi.core.extensions.fields import FieldsExtension
+from stac_fastapi.core.rate_limit import setup_rate_limit
 from stac_fastapi.core.route_dependencies import get_route_dependencies
 from stac_fastapi.core.session import Session
 from stac_fastapi.elasticsearch.config import ElasticsearchSettings
@@ -101,6 +101,7 @@
 # Add rate limit
 setup_rate_limit(app)
 
+
 @app.on_event("startup")
 async def _startup_event() -> None:
     await create_index_templates()

stac_fastapi/opensearch/stac_fastapi/opensearch/app.py

@@ -1,7 +1,6 @@
 """FastAPI application."""
 
 import os
-from stac_fastapi.core.rate_limit import setup_rate_limit
 
 from stac_fastapi.api.app import StacApi
 from stac_fastapi.api.models import create_get_request_model, create_post_request_model
@@ -18,6 +17,7 @@
     EsAsyncAggregationClient,
 )
 from stac_fastapi.core.extensions.fields import FieldsExtension
+from stac_fastapi.core.rate_limit import setup_rate_limit
 from stac_fastapi.core.route_dependencies import get_route_dependencies
 from stac_fastapi.core.session import Session
 from stac_fastapi.extensions.core import (
@@ -101,6 +101,7 @@
 # Add rate limit
 setup_rate_limit(app)
 
+
 @app.on_event("startup")
 async def _startup_event() -> None:
     await create_index_templates()

stac_fastapi/tests/conftest.py

@@ -24,8 +24,8 @@
     EsAggregationExtensionPostRequest,
     EsAsyncAggregationClient,
 )
-from stac_fastapi.core.route_dependencies import get_route_dependencies
 from stac_fastapi.core.rate_limit import setup_rate_limit
+from stac_fastapi.core.route_dependencies import get_route_dependencies
 
 if os.getenv("BACKEND", "elasticsearch").lower() == "opensearch":
     from stac_fastapi.opensearch.config import AsyncOpensearchSettings as AsyncSettings
@@ -243,9 +243,9 @@ async def app():
 @pytest_asyncio.fixture(scope="function")
 async def app_rate_limit(monkeypatch):
     monkeypatch.setenv("STAC_FASTAPI_RATE_LIMIT", "2/minute")
-    
+
     settings = AsyncSettings()
-    
+
     aggregation_extension = AggregationExtension(
         client=EsAsyncAggregationClient(
             database=database, session=None, settings=settings
@@ -292,7 +292,6 @@ async def app_rate_limit(monkeypatch):
     return app
 
 
-
 @pytest_asyncio.fixture(scope="session")
 async def app_client(app):
     await create_index_templates()

stac_fastapi/tests/rate_limit/test_rate_limit.py

@@ -1,20 +1,24 @@
+import logging
+
 import pytest
 from httpx import AsyncClient
-import logging
 from slowapi.errors import RateLimitExceeded
 
 logger = logging.getLogger(__name__)
 
+
 @pytest.mark.asyncio
 async def test_rate_limit(app_client_rate_limit: AsyncClient):
     expected_status_codes = [200, 200, 429, 429, 429]
-    
+
     for i, expected_status_code in enumerate(expected_status_codes):
         try:
             response = await app_client_rate_limit.get("/collections")
             status_code = response.status_code
         except RateLimitExceeded:
             status_code = 429
-        
+
         logger.info(f"Request {i+1}: Status code {status_code}")
-        assert status_code == expected_status_code, f"Expected status code {expected_status_code}, but got {status_code}"
+        assert (
+            status_code == expected_status_code
+        ), f"Expected status code {expected_status_code}, but got {status_code}"