diff --git a/CHANGELOG.md b/CHANGELOG.md index 1407ce4..7bf7649 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,15 +4,38 @@ We follow [Semantic Versions](https://semver.org/). ## WIP -## Version 0.2.5 +## Version 0.3.0 ### Features +- Introduced a custom error class `RatelimitIOError`, designed to mirror the behavior of `HTTPException`, making it easier to integrate with web frameworks like Flask, Django, and FastAPI. +- Enhanced decorator functionality + - Automatically applies `base_limit` when the decorator is used without explicit arguments. +- Added seamless support for distinguishing between incoming and outgoing requests using the `is_incoming` flag. +- Implemented 429 Too Many Requests handling: + - Provides examples for integrating with Flask, Django, and FastAPI to manage error responses gracefully. + +### Fixes + +- Fixed an issue where `unique_key` fallback to `default_key` was inconsistent. +- Improved Redis script loading to ensure proper error handling and robustness during high-concurrency operations. + +### Misc + +- Updated README with examples for: + - Using decorators and context managers for both sync and async workflows. + - Handling `429 Too Many Requests` errors in popular frameworks. +- Enhanced test coverage to verify all added features and scenarios. + +## Version 0.2.5 + +### Misc + - Adds shields `downloads` ## Version 0.2.4 -### Features +### Misc - Edits README.md @@ -36,7 +59,7 @@ We follow [Semantic Versions](https://semver.org/). ## Version 0.2.0 -### Features +### Misc - Adds wheel builds to GitHub Actions diff --git a/README.md b/README.md index 0ca4d9a..bb9dd79 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # RatelimitIO -A Python library for rate limiting, built to handle both incoming and outgoing requests efficiently. Supports both synchronous and asynchronous paradigms. Powered by Redis, this library provides decorators, context managers, and easy integration with APIs to manage request limits with precision. +A Python library for rate limiting, designed to handle both incoming and outgoing requests efficiently. Supports both synchronous and asynchronous paradigms. Powered by Redis, this library provides decorators, context managers, and easy integration with APIs to manage request limits with precision. #### Project Information [![Tests & Lint](https://github.com/bagowix/ratelimit-io/actions/workflows/actions.yml/badge.svg)](https://github.com/bagowix/ratelimit-io/actions/workflows/actions.yml) @@ -13,7 +13,7 @@ A Python library for rate limiting, built to handle both incoming and outgoing r ## Features -- **Incoming and Outgoing Support**: Effectively handles limits for both inbound +- **Incoming and Outgoing Support**: Effectively handles limits for both inbound (e.g., API requests) and outbound (e.g., client requests to external APIs) traffic. - **Synchronous and Asynchronous Support**: Works seamlessly in both paradigms. - **Redis Backend**: Leverages Redis for fast and scalable rate limiting. - **Flexible API**: @@ -22,6 +22,8 @@ A Python library for rate limiting, built to handle both incoming and outgoing r - Integrate directly into API clients, middlewares, or custom request handlers. - **Customizable Rate Limits**: Specify limits per key, time period, and requests. - **Robust Lua Script**: Ensures efficient and atomic rate limiting logic for high-concurrency use cases. +- **Automatic Error Handling**: Easily manage 429 Too Many Requests errors in popular frameworks like Flask, Django, and FastAPI. +- **Support for Incoming Request Behavior**: Use the `is_incoming` flag to distinguish between incoming requests (throwing errors immediately) and outgoing requests (waiting for available slots). - **Ease of Use**: Simple and intuitive integration into Python applications. --- @@ -43,9 +45,9 @@ from ratelimit_io import RatelimitIO, LimitSpec from redis import Redis redis_client = Redis(host="localhost", port=6379) -limiter = RatelimitIO(backend=redis_client) +limiter = RatelimitIO(backend=redis_client, base_limit=LimitSpec(requests=10, seconds=60)) -@limiter(LimitSpec(requests=10, seconds=60), unique_key="user:123") # unique_key is optional +@limiter def fetch_data(): return "Request succeeded!" @@ -53,8 +55,6 @@ def fetch_data(): fetch_data() ``` -**Note**: The `unique_key` parameter is `optional`. If not provided, the `IP address` (or another default identifier, based on your application logic) will be used to apply rate limiting. - ### Using as a Asynchronous Decorator ```python @@ -62,9 +62,9 @@ from ratelimit_io import RatelimitIO, LimitSpec from redis.asyncio import Redis as AsyncRedis async_redis_client = AsyncRedis(host="localhost", port=6379) -async_limiter = RatelimitIO(backend=async_redis_client) +async_limiter = RatelimitIO(backend=async_redis_client, base_limit=LimitSpec(requests=10, seconds=60)) -@async_limiter(LimitSpec(requests=10, seconds=60), unique_key="user:123") # unique_key is optional +@async_limiter async def fetch_data(): return "Request succeeded!" @@ -72,11 +72,36 @@ async def fetch_data(): await fetch_data() ``` -**Note**: Similar to the synchronous decorator, the `unique_key` is `optional`. If omitted, the `IP address` (or a default identifier) will be used. +### Incoming vs. Outgoing Request Handling (`is_incoming`) + +- When `is_incoming=True`, the rate limiter will immediately raise a `RatelimitIOError` when limits are exceeded. +- When `is_incoming=False` (default), the rate limiter will wait until a slot becomes available. + +```python +# Incoming request example (throws an error on limit breach) +limiter = RatelimitIO(backend=redis_client, is_incoming=True) + +@limiter(LimitSpec(requests=5, seconds=10)) +def fetch_data(): + return "Request succeeded!" + +# Outgoing request example (waits if limits are exceeded) +outgoing_limiter = RatelimitIO(backend=redis_client) + +@outgoing_limiter(LimitSpec(requests=5, seconds=10)) +def fetch_data_outgoing(): + return "Request succeeded!" +``` ### Using as a Synchronous Context Manager ```python +from ratelimit_io import RatelimitIO, LimitSpec +from redis import Redis + +redis_client = Redis(host="localhost", port=6379) +limiter = RatelimitIO(backend=redis_client) + with limiter: limiter.wait("user:456", LimitSpec(requests=5, seconds=10)) ``` @@ -84,8 +109,81 @@ with limiter: ### Using as a Asynchronous Context Manager ```python -async with limiter: - await limiter.a_wait("user:456", LimitSpec(requests=5, seconds=10)) +from ratelimit_io import RatelimitIO, LimitSpec +from redis.asyncio import Redis as AsyncRedis + +async_redis_client = AsyncRedis(host="localhost", port=6379) +async_limiter = RatelimitIO(backend=async_redis_client) + +async with async_limiter: + await async_limiter.a_wait("user:456", LimitSpec(requests=5, seconds=10)) +``` + +## Error Handling for 429 Responses + +### FastAPI Example + +```python +from fastapi import FastAPI, HTTPException, Request +from ratelimit_io import RatelimitIO, RatelimitIOError, LimitSpec +from redis.asyncio import Redis as AsyncRedis + +app = FastAPI() +redis_client = AsyncRedis(host="localhost", port=6379) +limiter = RatelimitIO(backend=redis_client, is_incoming=True) + +@app.middleware("http") +async def ratelimit_middleware(request: Request, call_next): + try: + await limiter.a_wait(f"user:{request.client.host}", LimitSpec(10, seconds=60)) + except RatelimitIOError as e: + raise HTTPException(status_code=e.status_code, detail=e.detail) + return await call_next(request) + +@app.get("/fetch") +async def fetch_data(): + return {"message": "Request succeeded!"} +``` + +### Django Middleware Example + +```python +from django.http import JsonResponse +from ratelimit_io import RatelimitIO, RatelimitIOError, LimitSpec +from redis import Redis + +redis_client = Redis(host="localhost", port=6379) +limiter = RatelimitIO(backend=redis_client, is_incoming=True) + +def ratelimit_middleware(get_response): + def middleware(request): + try: + limiter.wait(f"user:{request.META['REMOTE_ADDR']}", LimitSpec(10, seconds=60)) + except RatelimitIOError as e: + return JsonResponse({"error": e.detail}, status=e.status_code) + return get_response(request) + return middleware +``` + +### Flask Example + +```python +from flask import Flask, jsonify +from ratelimit_io import RatelimitIO, RatelimitIOError, LimitSpec +from redis import Redis + +app = Flask(__name__) +redis_client = Redis(host="localhost", port=6379) +limiter = RatelimitIO(backend=redis_client, is_incoming=True) + +@app.errorhandler(RatelimitIOError) +def handle_ratelimit_error(error): + return jsonify({"error": error.detail}), error.status_code + +@app.route("/fetch") +@limiter +def fetch_data(): + return jsonify({"message": "Request succeeded!"}) ``` ## License diff --git a/coverage.json b/coverage.json index 3735a0e..b4d7872 100644 --- a/coverage.json +++ b/coverage.json @@ -1 +1 @@ -{"meta": {"format": 3, "version": "7.6.1", "timestamp": "2025-01-14T21:12:18.195123", "branch_coverage": false, "show_contexts": false}, "files": {"ratelimit_io/__init__.py": {"executed_lines": [1, 3, 4, 5, 7, 10, 11, 14, 15, 25, 27, 28, 30, 31, 32, 33, 35, 36, 38, 45, 46, 47, 48, 49, 50, 51, 52, 54, 55, 58], "summary": {"covered_lines": 28, "num_statements": 28, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": [], "functions": {"LimitSpec.__init__": {"executed_lines": [27, 28, 30, 31, 32, 33, 35, 36], "summary": {"covered_lines": 8, "num_statements": 8, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec.total_seconds": {"executed_lines": [45, 46, 47, 48, 49, 50, 51, 52], "summary": {"covered_lines": 8, "num_statements": 8, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec.__str__": {"executed_lines": [55], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "": {"executed_lines": [1, 3, 4, 5, 7, 10, 11, 14, 15, 25, 38, 54, 58], "summary": {"covered_lines": 11, "num_statements": 11, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}, "classes": {"RatelimitIOError": {"executed_lines": [], "summary": {"covered_lines": 0, "num_statements": 0, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec": {"executed_lines": [27, 28, 30, 31, 32, 33, 35, 36, 45, 46, 47, 48, 49, 50, 51, 52, 55], "summary": {"covered_lines": 17, "num_statements": 17, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "": {"executed_lines": [1, 3, 4, 5, 7, 10, 11, 14, 15, 25, 38, 54, 58], "summary": {"covered_lines": 11, "num_statements": 11, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}}, "ratelimit_io/rate_limit.py": {"executed_lines": [1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 15, 16, 26, 39, 55, 59, 62, 63, 65, 74, 75, 77, 78, 79, 81, 96, 97, 99, 110, 111, 112, 113, 114, 118, 120, 121, 122, 123, 127, 129, 132, 134, 136, 137, 139, 141, 143, 145, 146, 148, 150, 152, 172, 174, 175, 176, 178, 179, 182, 183, 185, 202, 204, 205, 206, 208, 209, 212, 213, 215, 229, 230, 232, 237, 243, 251, 253, 255, 257, 268, 269, 276, 277, 278, 280, 281, 282, 283, 290, 291, 299, 311, 312, 319, 320, 322, 323, 324, 331, 336, 346, 348, 350, 351, 352, 354, 356, 357, 358], "summary": {"covered_lines": 106, "num_statements": 137, "percent_covered": 77.37226277372262, "percent_covered_display": "77", "missing_lines": 31, "excluded_lines": 0}, "missing_lines": [28, 29, 31, 32, 33, 34, 36, 37, 46, 47, 48, 49, 50, 51, 52, 53, 56, 115, 116, 124, 125, 180, 210, 254, 292, 294, 295, 296, 297, 332, 334], "excluded_lines": [], "functions": {"LimitSpec.__init__": {"executed_lines": [], "summary": {"covered_lines": 0, "num_statements": 8, "percent_covered": 0.0, "percent_covered_display": "0", "missing_lines": 8, "excluded_lines": 0}, "missing_lines": [28, 29, 31, 32, 33, 34, 36, 37], "excluded_lines": []}, "LimitSpec.total_seconds": {"executed_lines": [], "summary": {"covered_lines": 0, "num_statements": 8, "percent_covered": 0.0, "percent_covered_display": "0", "missing_lines": 8, "excluded_lines": 0}, "missing_lines": [46, 47, 48, 49, 50, 51, 52, 53], "excluded_lines": []}, "LimitSpec.__str__": {"executed_lines": [], "summary": {"covered_lines": 0, "num_statements": 1, "percent_covered": 0.0, "percent_covered_display": "0", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [56], "excluded_lines": []}, "RatelimitIO.__init__": {"executed_lines": [74, 75, 77, 78, 79, 81, 96, 97], "summary": {"covered_lines": 8, "num_statements": 8, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__call__": {"executed_lines": [110, 132], "summary": {"covered_lines": 2, "num_statements": 2, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__call__.decorator": {"executed_lines": [111, 120, 129], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__call__.decorator.async_wrapper": {"executed_lines": [112, 113, 114, 118], "summary": {"covered_lines": 4, "num_statements": 6, "percent_covered": 66.66666666666667, "percent_covered_display": "67", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [115, 116], "excluded_lines": []}, "RatelimitIO.__call__.decorator.sync_wrapper": {"executed_lines": [121, 122, 123, 127], "summary": {"covered_lines": 4, "num_statements": 6, "percent_covered": 66.66666666666667, "percent_covered_display": "67", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [124, 125], "excluded_lines": []}, "RatelimitIO.__enter__": {"executed_lines": [136, 137], "summary": {"covered_lines": 2, "num_statements": 2, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__exit__": {"executed_lines": [141], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__aenter__": {"executed_lines": [145, 146], "summary": {"covered_lines": 2, "num_statements": 2, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__aexit__": {"executed_lines": [150], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.wait": {"executed_lines": [172, 174, 175, 176, 178, 179, 182, 183], "summary": {"covered_lines": 8, "num_statements": 9, "percent_covered": 88.88888888888889, "percent_covered_display": "89", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [180], "excluded_lines": []}, "RatelimitIO.a_wait": {"executed_lines": [202, 204, 205, 206, 208, 209, 212, 213], "summary": {"covered_lines": 8, "num_statements": 9, "percent_covered": 88.88888888888889, "percent_covered_display": "89", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [210], "excluded_lines": []}, "RatelimitIO._prepare_key_and_limit": {"executed_lines": [229, 230, 232, 237, 243, 251, 253, 255], "summary": {"covered_lines": 8, "num_statements": 9, "percent_covered": 88.88888888888889, "percent_covered_display": "89", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [254], "excluded_lines": []}, "RatelimitIO._enforce_limit_sync": {"executed_lines": [268, 269, 276, 277, 278, 280, 281, 282, 283, 290, 291], "summary": {"covered_lines": 11, "num_statements": 16, "percent_covered": 68.75, "percent_covered_display": "69", "missing_lines": 5, "excluded_lines": 0}, "missing_lines": [292, 294, 295, 296, 297], "excluded_lines": []}, "RatelimitIO._enforce_limit_async": {"executed_lines": [311, 312, 319, 320, 322, 323, 324, 331], "summary": {"covered_lines": 8, "num_statements": 10, "percent_covered": 80.0, "percent_covered_display": "80", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [332, 334], "excluded_lines": []}, "RatelimitIO._generate_key": {"executed_lines": [346], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO._ensure_script_loaded_sync": {"executed_lines": [350, 351, 352], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO._ensure_script_loaded_async": {"executed_lines": [356, 357, 358], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "": {"executed_lines": [1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 15, 16, 26, 39, 55, 59, 62, 63, 65, 99, 134, 139, 143, 148, 152, 185, 215, 257, 299, 336, 348, 354], "summary": {"covered_lines": 29, "num_statements": 29, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}, "classes": {"RatelimitIOError": {"executed_lines": [], "summary": {"covered_lines": 0, "num_statements": 0, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec": {"executed_lines": [], "summary": {"covered_lines": 0, "num_statements": 17, "percent_covered": 0.0, "percent_covered_display": "0", "missing_lines": 17, "excluded_lines": 0}, "missing_lines": [28, 29, 31, 32, 33, 34, 36, 37, 46, 47, 48, 49, 50, 51, 52, 53, 56], "excluded_lines": []}, "RatelimitIO": {"executed_lines": [74, 75, 77, 78, 79, 81, 96, 97, 110, 111, 112, 113, 114, 118, 120, 121, 122, 123, 127, 129, 132, 136, 137, 141, 145, 146, 150, 172, 174, 175, 176, 178, 179, 182, 183, 202, 204, 205, 206, 208, 209, 212, 213, 229, 230, 232, 237, 243, 251, 253, 255, 268, 269, 276, 277, 278, 280, 281, 282, 283, 290, 291, 311, 312, 319, 320, 322, 323, 324, 331, 346, 350, 351, 352, 356, 357, 358], "summary": {"covered_lines": 77, "num_statements": 91, "percent_covered": 84.61538461538461, "percent_covered_display": "85", "missing_lines": 14, "excluded_lines": 0}, "missing_lines": [115, 116, 124, 125, 180, 210, 254, 292, 294, 295, 296, 297, 332, 334], "excluded_lines": []}, "": {"executed_lines": [1, 2, 3, 4, 5, 6, 7, 8, 11, 12, 15, 16, 26, 39, 55, 59, 62, 63, 65, 99, 134, 139, 143, 148, 152, 185, 215, 257, 299, 336, 348, 354], "summary": {"covered_lines": 29, "num_statements": 29, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}}}, "totals": {"covered_lines": 134, "num_statements": 165, "percent_covered": 81.21212121212122, "percent_covered_display": "81", "missing_lines": 31, "excluded_lines": 0}} +{"meta": {"format": 3, "version": "7.6.1", "timestamp": "2025-01-15T13:36:25.642529", "branch_coverage": false, "show_contexts": false}, "files": {"ratelimit_io/__init__.py": {"executed_lines": [1, 2, 3, 5], "summary": {"covered_lines": 4, "num_statements": 4, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": [], "functions": {"": {"executed_lines": [1, 2, 3, 5], "summary": {"covered_lines": 4, "num_statements": 4, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}, "classes": {"": {"executed_lines": [1, 2, 3, 5], "summary": {"covered_lines": 4, "num_statements": 4, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}}, "ratelimit_io/rate_limit.py": {"executed_lines": [1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 16, 17, 19, 24, 25, 26, 29, 30, 40, 47, 48, 50, 51, 52, 53, 55, 56, 61, 68, 69, 70, 71, 72, 73, 74, 75, 77, 78, 81, 84, 85, 87, 108, 109, 111, 112, 113, 114, 115, 117, 132, 133, 135, 155, 156, 158, 159, 165, 166, 167, 168, 173, 180, 181, 182, 183, 184, 188, 190, 191, 192, 197, 204, 205, 206, 207, 208, 212, 214, 220, 222, 224, 225, 227, 229, 231, 233, 234, 236, 238, 240, 265, 267, 269, 270, 271, 272, 274, 275, 277, 278, 282, 283, 285, 310, 312, 314, 315, 316, 317, 319, 320, 322, 323, 327, 328, 330, 347, 348, 350, 355, 361, 366, 368, 370, 372, 383, 384, 391, 392, 393, 394, 395, 402, 408, 421, 422, 429, 430, 431, 432, 433, 440, 446, 456, 458, 460, 463, 464, 466, 468, 469, 470], "summary": {"covered_lines": 152, "num_statements": 164, "percent_covered": 92.6829268292683, "percent_covered_display": "93", "missing_lines": 12, "excluded_lines": 0}, "missing_lines": [160, 174, 185, 198, 209, 279, 324, 369, 403, 404, 441, 442], "excluded_lines": [], "functions": {"RatelimitIOError.__init__": {"executed_lines": [24, 25, 26], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec.__init__": {"executed_lines": [47, 48, 50, 51, 52, 53, 55, 56], "summary": {"covered_lines": 8, "num_statements": 8, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec.total_seconds": {"executed_lines": [68, 69, 70, 71, 72, 73, 74, 75], "summary": {"covered_lines": 8, "num_statements": 8, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec.__str__": {"executed_lines": [78], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__init__": {"executed_lines": [108, 109, 111, 112, 113, 114, 115, 117, 132, 133], "summary": {"covered_lines": 10, "num_statements": 10, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__call__": {"executed_lines": [155, 156, 158, 159, 165, 166, 220], "summary": {"covered_lines": 7, "num_statements": 8, "percent_covered": 87.5, "percent_covered_display": "88", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [160], "excluded_lines": []}, "RatelimitIO.__call__.decorator": {"executed_lines": [167, 190, 191, 214], "summary": {"covered_lines": 4, "num_statements": 4, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__call__.decorator.async_wrapper": {"executed_lines": [168, 173, 180, 181, 182, 183, 184, 188], "summary": {"covered_lines": 8, "num_statements": 10, "percent_covered": 80.0, "percent_covered_display": "80", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [174, 185], "excluded_lines": []}, "RatelimitIO.__call__.decorator.sync_wrapper": {"executed_lines": [192, 197, 204, 205, 206, 207, 208, 212], "summary": {"covered_lines": 8, "num_statements": 10, "percent_covered": 80.0, "percent_covered_display": "80", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [198, 209], "excluded_lines": []}, "RatelimitIO.__enter__": {"executed_lines": [224, 225], "summary": {"covered_lines": 2, "num_statements": 2, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__exit__": {"executed_lines": [229], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__aenter__": {"executed_lines": [233, 234], "summary": {"covered_lines": 2, "num_statements": 2, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.__aexit__": {"executed_lines": [238], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO.wait": {"executed_lines": [265, 267, 269, 270, 271, 272, 274, 275, 277, 278, 282, 283], "summary": {"covered_lines": 12, "num_statements": 13, "percent_covered": 92.3076923076923, "percent_covered_display": "92", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [279], "excluded_lines": []}, "RatelimitIO.a_wait": {"executed_lines": [310, 312, 314, 315, 316, 317, 319, 320, 322, 323, 327, 328], "summary": {"covered_lines": 12, "num_statements": 13, "percent_covered": 92.3076923076923, "percent_covered_display": "92", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [324], "excluded_lines": []}, "RatelimitIO._prepare_key_and_limit": {"executed_lines": [347, 348, 350, 355, 361, 366, 368, 370], "summary": {"covered_lines": 8, "num_statements": 9, "percent_covered": 88.88888888888889, "percent_covered_display": "89", "missing_lines": 1, "excluded_lines": 0}, "missing_lines": [369], "excluded_lines": []}, "RatelimitIO._enforce_limit_sync": {"executed_lines": [383, 384, 391, 392, 393, 394, 395, 402], "summary": {"covered_lines": 8, "num_statements": 10, "percent_covered": 80.0, "percent_covered_display": "80", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [403, 404], "excluded_lines": []}, "RatelimitIO._enforce_limit_async": {"executed_lines": [421, 422, 429, 430, 431, 432, 433, 440], "summary": {"covered_lines": 8, "num_statements": 10, "percent_covered": 80.0, "percent_covered_display": "80", "missing_lines": 2, "excluded_lines": 0}, "missing_lines": [441, 442], "excluded_lines": []}, "RatelimitIO._generate_key": {"executed_lines": [456], "summary": {"covered_lines": 1, "num_statements": 1, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO._ensure_script_loaded_sync": {"executed_lines": [460, 463, 464], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO._ensure_script_loaded_async": {"executed_lines": [468, 469, 470], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "": {"executed_lines": [1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 16, 17, 19, 29, 30, 40, 61, 77, 81, 84, 85, 87, 135, 222, 227, 231, 236, 240, 285, 330, 372, 408, 446, 458, 466], "summary": {"covered_lines": 34, "num_statements": 34, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}, "classes": {"RatelimitIOError": {"executed_lines": [24, 25, 26], "summary": {"covered_lines": 3, "num_statements": 3, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "LimitSpec": {"executed_lines": [47, 48, 50, 51, 52, 53, 55, 56, 68, 69, 70, 71, 72, 73, 74, 75, 78], "summary": {"covered_lines": 17, "num_statements": 17, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}, "RatelimitIO": {"executed_lines": [108, 109, 111, 112, 113, 114, 115, 117, 132, 133, 155, 156, 158, 159, 165, 166, 167, 168, 173, 180, 181, 182, 183, 184, 188, 190, 191, 192, 197, 204, 205, 206, 207, 208, 212, 214, 220, 224, 225, 229, 233, 234, 238, 265, 267, 269, 270, 271, 272, 274, 275, 277, 278, 282, 283, 310, 312, 314, 315, 316, 317, 319, 320, 322, 323, 327, 328, 347, 348, 350, 355, 361, 366, 368, 370, 383, 384, 391, 392, 393, 394, 395, 402, 421, 422, 429, 430, 431, 432, 433, 440, 456, 460, 463, 464, 468, 469, 470], "summary": {"covered_lines": 98, "num_statements": 110, "percent_covered": 89.0909090909091, "percent_covered_display": "89", "missing_lines": 12, "excluded_lines": 0}, "missing_lines": [160, 174, 185, 198, 209, 279, 324, 369, 403, 404, 441, 442], "excluded_lines": []}, "": {"executed_lines": [1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 16, 17, 19, 29, 30, 40, 61, 77, 81, 84, 85, 87, 135, 222, 227, 231, 236, 240, 285, 330, 372, 408, 446, 458, 466], "summary": {"covered_lines": 34, "num_statements": 34, "percent_covered": 100.0, "percent_covered_display": "100", "missing_lines": 0, "excluded_lines": 0}, "missing_lines": [], "excluded_lines": []}}}}, "totals": {"covered_lines": 156, "num_statements": 168, "percent_covered": 92.85714285714286, "percent_covered_display": "93", "missing_lines": 12, "excluded_lines": 0}} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 556819b..be2668f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "ratelimit-io" -version = "0.2.5" +version = "0.3.0" description = "Flexible bidirectional rate-limiting library with redis backend" authors = ["Galushko Bogdan "] license = "MIT" diff --git a/ratelimit_io/__init__.py b/ratelimit_io/__init__.py index 49d4433..8a852e7 100644 --- a/ratelimit_io/__init__.py +++ b/ratelimit_io/__init__.py @@ -1,67 +1,9 @@ -from typing import Optional -from typing import Union - -from redis import Redis -from redis.asyncio import Redis as AsyncRedis -from typing_extensions import TypeAlias - -from ratelimit_io.rate_limit import RatelimitIO # noqa - - -class RatelimitIOError(Exception): - """Raised when the rate limit is exceeded.""" - - -class LimitSpec: - """ - Specifies the number of requests allowed in a time frame. - - Attributes: - requests (int): Maximum number of requests. - seconds (Optional[int]): Time frame in seconds. - minutes (Optional[int]): Time frame in minutes. - hours (Optional[int]): Time frame in hours. - """ - - def __init__( - self, - requests: int, - seconds: Optional[int] = None, - minutes: Optional[int] = None, - hours: Optional[int] = None, - ): - if requests <= 0: - raise ValueError("Requests must be greater than 0.") - - self.requests = requests - self.seconds = seconds - self.minutes = minutes - self.hours = hours - - if self.total_seconds() == 0: - raise ValueError( - "At least one time frame " - "(seconds, minutes, or hours) must be provided." - ) - - def total_seconds(self) -> int: - """ - Calculates the total time frame in seconds. - - Returns: - int: Total time in seconds. - """ - total = 0 - if self.seconds: - total += self.seconds - if self.minutes: - total += self.minutes * 60 - if self.hours: - total += self.hours * 3600 - return total - - def __str__(self) -> str: - return f"{self.requests}/{self.total_seconds()}s" - - -RedisBackend: TypeAlias = Union[Redis, AsyncRedis] +from ratelimit_io.rate_limit import LimitSpec +from ratelimit_io.rate_limit import RatelimitIO +from ratelimit_io.rate_limit import RatelimitIOError + +__all__ = [ + "RatelimitIO", + "LimitSpec", + "RatelimitIOError", +] diff --git a/ratelimit_io/rate_limit.py b/ratelimit_io/rate_limit.py index 2e03c12..28a1dd4 100644 --- a/ratelimit_io/rate_limit.py +++ b/ratelimit_io/rate_limit.py @@ -1,4 +1,5 @@ import time +from functools import wraps from typing import Callable from typing import Optional from typing import Tuple @@ -15,6 +16,15 @@ class RatelimitIOError(Exception): """Raised when the rate limit is exceeded.""" + def __init__( + self, + detail: Optional[str] = "Too many Requests", + status_code: Optional[int] = 429, + ) -> None: + super().__init__(detail) + self.status_code = status_code + self.detail = detail + class LimitSpec: """ @@ -33,7 +43,7 @@ def __init__( seconds: Optional[int] = None, minutes: Optional[int] = None, hours: Optional[int] = None, - ): + ) -> None: if requests <= 0: raise ValueError("Requests must be greater than 0.") @@ -77,24 +87,32 @@ class RatelimitIO: def __init__( self, backend: RedisBackend, + is_incoming: bool = False, base_url: Optional[str] = None, base_limit: Optional[LimitSpec] = None, + default_key: Optional[str] = "unknown_key", ): """ Initializes the RatelimitIO instance. Args: backend (Redis | AsyncRedis): Redis backend instance. + is_incoming (bool): Whether the rate limiter is + for incoming requests. base_url (Optional[str]): Base URL for outgoing request limits. base_limit (Optional[LimitSpec]): Default rate limit for the base URL. + default_key (Optional[str]): Default unique key for rate limiting. + Defaults to "unknown_key". """ if not isinstance(backend, (Redis, AsyncRedis)): raise RuntimeError("Unsupported Redis backend.") self.backend = backend + self.is_incoming = is_incoming self.base_url = base_url self.base_limit = base_limit + self.default_key = default_key self._lua_script = b""" local current = redis.call("GET", KEYS[1]) @@ -115,38 +133,82 @@ def __init__( self._script_loaded = False def __call__( - self, limit_spec: LimitSpec, unique_key: Optional[str] = None + self, + func: Optional[Callable] = None, + *, + limit_spec: Optional[LimitSpec] = None, + unique_key: Optional[str] = None, ) -> Callable: """ Decorator for applying rate limits to functions. Args: - limit_spec (LimitSpec): Rate limit specification. + func (Callable): Function to decorate. + limit_spec (Optional[LimitSpec]): Rate limit specification. + Defaults to `self.base_limit`. unique_key (Optional[str]): Optional unique key for rate limiting. + If not provided, tries `self.default_key` or `kwargs["ip"]`. Returns: Callable: Decorated function. """ + if func and callable(func): + return self(limit_spec=self.base_limit)(func) + + limit_spec = limit_spec or self.base_limit + if not limit_spec: + raise ValueError( + "Rate limit specification is missing. Provide a limit_spec " + "or ensure base_limit is set during initialization." + ) def decorator(func: Callable) -> Callable: + @wraps(func) async def async_wrapper(*args, **kwargs): - key = unique_key or kwargs.get("ip", "unknown_ip") + key = ( + unique_key + or self.default_key + or kwargs.get("ip", "unknown_ip") + ) + if not key: + raise ValueError( + "Unique key is required. Provide `unique_key`, " + "set `default_key`, " + "or include `ip` in kwargs." + ) + try: - await self.a_wait(f"incoming:{key}", limit_spec) - except RatelimitIOError as exc: + await self.a_wait(f"ratelimit:{key}", limit_spec) + except RatelimitIOError: + if self.is_incoming: + raise raise RuntimeError( - f"Rate limit exceeded in {func.__name__}: {str(exc)}" - ) from exc + f"Rate limit exceeded in {func.__name__}" + ) from None return await func(*args, **kwargs) + @wraps(func) def sync_wrapper(*args, **kwargs): - key = unique_key or kwargs.get("ip", "unknown_ip") + key = ( + unique_key + or self.default_key + or kwargs.get("ip", "unknown_ip") + ) + if not key: + raise ValueError( + "Unique key is required. Provide `unique_key`, " + "set `default_key`, " + "or include `ip` in kwargs." + ) + try: - self.wait(f"incoming:{key}", limit_spec) - except RatelimitIOError as exc: + self.wait(f"ratelimit:{key}", limit_spec) + except RatelimitIOError: + if self.is_incoming: + raise raise RuntimeError( - f"Rate limit exceeded in {func.__name__}: {str(exc)}" - ) from exc + f"Rate limit exceeded in {func.__name__}" + ) from None return func(*args, **kwargs) return ( @@ -203,6 +265,12 @@ def wait( self._ensure_script_loaded_sync() key, limit_spec = self._prepare_key_and_limit(key, limit_spec) + + if self.is_incoming: + if not self._enforce_limit_sync(key, limit_spec): + raise RatelimitIOError() + return + start_time = time.time() backoff = backoff_start @@ -242,6 +310,12 @@ async def a_wait( await self._ensure_script_loaded_async() key, limit_spec = self._prepare_key_and_limit(key, limit_spec) + + if self.is_incoming: + if not await self._enforce_limit_async(key, limit_spec): + raise RatelimitIOError() + return + start_time = time.time() backoff = backoff_start @@ -354,7 +428,6 @@ async def _enforce_limit_async( ) return bool(allowed) except NoScriptError: - # Attempt to load the script and retry once await self._ensure_script_loaded_async() try: allowed = await self.backend.evalsha( # type: ignore @@ -366,7 +439,6 @@ async def _enforce_limit_async( ) return bool(allowed) except NoScriptError as exc: - # If script loading fails again, raise an error raise RuntimeError( "Failed to load Lua script into Redis." ) from exc diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py index 77ba3d1..e27f764 100644 --- a/tests/test_rate_limit.py +++ b/tests/test_rate_limit.py @@ -1,6 +1,8 @@ import os import re import time +from typing import AsyncGenerator +from typing import Generator import pytest import redis @@ -9,10 +11,11 @@ from ratelimit_io import LimitSpec from ratelimit_io import RatelimitIO +from ratelimit_io import RatelimitIOError @pytest.fixture(scope="function") -def real_redis_client(): +def real_redis_client() -> Generator[Redis, None, None]: """ Real Redis client for synchronous tests. @@ -33,7 +36,7 @@ def real_redis_client(): @pytest.fixture(scope="function") -async def real_async_redis_client(): +async def real_async_redis_client() -> AsyncGenerator[AsyncRedis, None]: """ Real Redis client for asynchronous tests. @@ -50,11 +53,11 @@ async def real_async_redis_client(): client = AsyncRedis(host=host, port=port, decode_responses=True) await client.flushall() yield client - await client.aclose() + await client.aclose() # type: ignore @pytest.fixture -def limiter(real_redis_client): +def limiter(real_redis_client) -> RatelimitIO: """Fixture for RatelimitIO instance with real Redis.""" return RatelimitIO( backend=real_redis_client, @@ -64,7 +67,7 @@ def limiter(real_redis_client): @pytest.fixture -async def async_limiter(real_async_redis_client): +async def async_limiter(real_async_redis_client) -> RatelimitIO: """Fixture for RatelimitIO instance with real Redis (async).""" return RatelimitIO( backend=real_async_redis_client, @@ -73,85 +76,167 @@ async def async_limiter(real_async_redis_client): ) -def test_sync_limit(limiter): - """Test synchronous rate limiting.""" - key = "sync_test" +def test_sync_limit_incoming(limiter): + """Test synchronous rate limiting with is_incoming=True.""" + limiter.is_incoming = True + key = "sync_incoming_test" for _ in range(5): limiter.wait(key, LimitSpec(requests=5, seconds=1)) - start_time = time.time() + with pytest.raises(RatelimitIOError, match="Too many Requests"): + limiter.wait(key, LimitSpec(requests=5, seconds=1)) - limiter.wait(key, LimitSpec(requests=5, seconds=1)) - elapsed_time = time.time() - start_time +@pytest.mark.asyncio +async def test_async_limit_incoming(async_limiter): + """Test asynchronous rate limiting with is_incoming=True.""" + async_limiter.is_incoming = True + key = "async_incoming_test" - assert elapsed_time >= 0.9, "Wait time did not occur as expected" + for _ in range(5): + await async_limiter.a_wait(key, LimitSpec(requests=5, seconds=1)) + + with pytest.raises(RatelimitIOError, match="Too many Requests"): + await async_limiter.a_wait(key, LimitSpec(requests=5, seconds=1)) @pytest.mark.asyncio -async def test_async_limit(async_limiter): - """Test asynchronous rate limiting.""" - key = "async_test" +async def test_async_decorator_incoming(async_limiter): + """Test asynchronous decorator usage with is_incoming=True.""" + async_limiter.is_incoming = True - await async_limiter.backend.flushall() + @async_limiter( + LimitSpec(requests=5, seconds=1), unique_key="async_test_key" + ) + async def limited_function(): + return "success" for _ in range(5): - await async_limiter.a_wait(key, LimitSpec(requests=5, seconds=5)) - await async_limiter.backend.get(async_limiter._generate_key(key)) + assert await limited_function() == "success" + + with pytest.raises(RatelimitIOError, match="Too many Requests"): + await limited_function() + + +def test_sync_decorator_incoming(limiter): + """Test synchronous decorator usage with is_incoming=True.""" + limiter.is_incoming = True + + @limiter(LimitSpec(requests=5, seconds=1), unique_key="sync_test_key") + def limited_function(): + return "success" + + for _ in range(5): + assert limited_function() == "success" + + with pytest.raises(RatelimitIOError, match="Too many Requests"): + limited_function() + + +def test_sync_limit_outgoing(limiter): + """Test synchronous rate limiting with is_incoming=False.""" + limiter.is_incoming = False + key = "sync_outgoing_test" + + for _ in range(5): + limiter.wait(key, LimitSpec(requests=5, seconds=1)) start_time = time.time() - await async_limiter.a_wait(key, LimitSpec(requests=5, seconds=5)) + limiter.wait(key, LimitSpec(requests=5, seconds=1)) elapsed_time = time.time() - start_time - assert elapsed_time >= 0.9, "Wait time did not occur as expected" - remaining = await async_limiter.backend.get( - async_limiter._generate_key(key) - ) - ttl = await async_limiter.backend.ttl(async_limiter._generate_key(key)) + assert elapsed_time >= 0.9, "Wait time not applied for outgoing request" - assert remaining is not None, "Key should exist in Redis" - assert int(remaining) <= 5, f"Unexpected requests count: {remaining}" - assert ttl > 0, f"Unexpected TTL: {ttl}" + +@pytest.mark.asyncio +async def test_async_limit_outgoing(async_limiter): + """Test asynchronous rate limiting with is_incoming=False.""" + async_limiter.is_incoming = False + key = "async_outgoing_test" + + for _ in range(5): + await async_limiter.a_wait(key, LimitSpec(requests=5, seconds=1)) + + start_time = time.time() + await async_limiter.a_wait(key, LimitSpec(requests=5, seconds=1)) + elapsed_time = time.time() - start_time + + assert elapsed_time >= 0.9, "Wait time not applied for outgoing request" + + +def test_default_key_incoming_behavior(limiter): + """Test default_key behavior when is_incoming=True.""" + limiter.is_incoming = True + limiter.default_key = "default_key_incoming" + + @limiter + def limited_function(): + return "success" + + for _ in range(5): + assert limited_function() == "success" + + with pytest.raises(RatelimitIOError, match="Too many Requests"): + limited_function() @pytest.mark.asyncio -async def test_async_decorator(async_limiter): - """Test asynchronous decorator usage.""" +async def test_default_key_incoming_behavior_async(async_limiter): + """Test default_key behavior when is_incoming=True (async).""" + async_limiter.is_incoming = True + async_limiter.default_key = "default_async_key" - @async_limiter( - LimitSpec(requests=5, seconds=1), unique_key="async_decorator_test" - ) + @async_limiter async def limited_function(): return "success" for _ in range(5): assert await limited_function() == "success" - start_time = time.time() - assert await limited_function() == "success" - elapsed_time = time.time() - start_time + with pytest.raises(RatelimitIOError, match="Too many Requests"): + await limited_function() - assert elapsed_time >= 0.9, "Wait time did not occur as expected" +def test_sync_decorator_without_args(limiter): + """Test synchronous decorator without arguments.""" -def test_sync_decorator(limiter): - """Test synchronous decorator usage.""" + limiter.default_key = "default_test_key" - @limiter( - LimitSpec(requests=5, seconds=1), unique_key="sync_decorator_test" - ) + @limiter def limited_function(): return "success" for _ in range(5): assert limited_function() == "success" + # Превышение лимита вызывает задержку + start_time = time.time() + limited_function() + elapsed_time = time.time() - start_time + + assert elapsed_time >= 0.9, "Wait time not applied with default_key" + + +@pytest.mark.asyncio +async def test_async_decorator_without_args(async_limiter): + """Test asynchronous decorator without arguments.""" + + async_limiter.default_key = "default_async_key" + + @async_limiter + async def limited_function(): + return "success" + + for _ in range(5): + assert await limited_function() == "success" + + # Превышение лимита вызывает задержку start_time = time.time() - assert limited_function() == "success" + await limited_function() elapsed_time = time.time() - start_time - assert elapsed_time >= 0.9, "Wait time did not occur as expected" + assert elapsed_time >= 0.9, "Wait time not applied with default_key" @pytest.mark.asyncio @@ -439,3 +524,104 @@ def test_limitspec_no_time_frame(): """Test LimitSpec raises an error when no time frame is provided.""" with pytest.raises(ValueError, match="At least one time frame"): LimitSpec(requests=5, seconds=0, minutes=0, hours=0) + + +def test_default_key_usage(limiter): + """Test the usage of default_key when unique_key is not provided.""" + limiter.default_key = "default_test_key" + + @limiter + def limited_function(): + return "success" + + # Ограничение должно работать с default_key + for _ in range(5): + assert limited_function() == "success" + + # Превышение лимита вызывает задержку + start_time = time.time() + limited_function() + elapsed_time = time.time() - start_time + + assert elapsed_time >= 0.9, "Wait time not applied with default_key" + + +def test_override_default_key(limiter): + """Test overriding default_key with unique_key.""" + limiter.default_key = "default_test_key" + + @limiter( + limit_spec=LimitSpec(requests=3, seconds=2), unique_key="custom_key" + ) + def limited_function(): + return "success" + + # Ограничение должно работать с unique_key + for _ in range(3): + assert limited_function() == "success" + + # Превышение лимита вызывает задержку + start_time = time.time() + limited_function() + elapsed_time = time.time() - start_time + + assert elapsed_time >= 1.9, "Wait time not applied with custom unique_key" + + +def test_priority_unique_key_over_default_key(limiter): + """Test that unique_key takes precedence over default_key.""" + + limiter.default_key = "default_test_key" + + @limiter( + limit_spec=LimitSpec(requests=5, seconds=1), unique_key="priority_key" + ) + def limited_function(): + return "success" + + for _ in range(5): + assert limited_function() == "success" + + start_time = time.time() + limited_function() + elapsed_time = time.time() - start_time + + assert elapsed_time >= 0.9, "Wait time not applied with unique_key" + + +def test_priority_default_key_over_ip(limiter): + """Test that default_key takes precedence over ip from kwargs.""" + + limiter.default_key = "default_test_key" + + @limiter + def limited_function(ip="user_ip"): + return "success" + + for _ in range(5): + assert limited_function() == "success" + + start_time = time.time() + limited_function(ip="ignored_ip") + elapsed_time = time.time() - start_time + + assert elapsed_time >= 0.9, "Wait time not applied with default_key" + + +def test_ip_key_as_fallback(limiter): + """Test that ip from kwargs is used if no other keys are provided.""" + + limiter.default_key = None + + @limiter + def limited_function(ip="user_ip"): + return "success" + + for _ in range(5): + assert limited_function() == "success" + + start_time = time.time() + limited_function() + elapsed_time = time.time() - start_time + + assert elapsed_time >= 0.9, "Wait time not applied with ip key"