diff --git a/rate_limiters/main.py b/rate_limiters/main.py deleted file mode 100644 index 72d791b..0000000 --- a/rate_limiters/main.py +++ /dev/null @@ -1,6 +0,0 @@ -def main(): - print("Hello from rate-limiters!") - - -if __name__ == "__main__": - main() diff --git a/rate_limiters/middleware_starlette.py b/rate_limiters/middleware_starlette.py index dfef511..e14ea5e 100644 --- a/rate_limiters/middleware_starlette.py +++ b/rate_limiters/middleware_starlette.py @@ -25,7 +25,6 @@ class RateLimiterMiddleware: if not limiter: limiter = self.clients.setdefault(client_ip, AsyncTokenBucketLimiter(10, 10)) - self.clients[client_ip] = limiter try: async with limiter: @@ -48,4 +47,6 @@ routes = [Route("/", main)] middleware = [Middleware(RateLimiterMiddleware)] -app = Starlette(routes=routes, middleware=middleware) + +def create_app() -> Starlette: + return Starlette(routes=routes, middleware=middleware) diff --git a/rate_limiters/start_app.sh b/rate_limiters/start_app.sh index e09c3ee..1b12c8a 100755 --- a/rate_limiters/start_app.sh +++ b/rate_limiters/start_app.sh @@ -1,3 +1,3 @@ #!/bin/bash -uvicorn --host 0.0.0.0 --port 8001 --reload middleware_starlette:app \ No newline at end of file +uvicorn --host 0.0.0.0 --port 8001 --reload --factory middleware_starlette:create_app \ No newline at end of file diff --git a/rate_limiters/tests/test_middleware_starlette.py b/rate_limiters/tests/test_middleware_starlette.py index e52c5a2..04818de 100644 --- a/rate_limiters/tests/test_middleware_starlette.py +++ b/rate_limiters/tests/test_middleware_starlette.py @@ -1,9 +1,8 @@ from time import sleep -from starlette.applications import Starlette from starlette.testclient import TestClient -from middleware_starlette import app, RateLimiterMiddleware, routes +from middleware_starlette import app def test_app_runs(): client = TestClient(app)