Skip to content

Commit

Permalink
Add unit test for bucket_name and fix associated issues
Browse files Browse the repository at this point in the history
  • Loading branch information
Aaron Fulton committed Feb 29, 2024
1 parent 6c9764d commit cf6d404
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 4 deletions.
3 changes: 3 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,9 @@ session.get('https://api.some_site.com/v1/users/1234')
For advanced use cases, you can define your own custom tracking behavior with the `bucket` option.
For example, an API that enforces rate limits based on a tenant ID, this feature can be used to track
rate limits per tenant. If `bucket` is specified, host tracking is disabled.

Note: It is advisable to use SQLite or Redis backends when using custom tracking because using the default backend
each session will track rate limits independently, even if both sessions call the same URL.
```python
sessionA = LimiterSession(per_second=5, bucket='tenant1')
sessionB = LimiterSession(per_second=5, bucket='tenant2')
Expand Down
7 changes: 3 additions & 4 deletions requests_ratelimiter/requests_ratelimiter.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ def __init__(
max_delay: Union[int, float, None] = None,
per_host: bool = True,
limit_statuses: Iterable[int] = (429,),
bucket: Optional[str] = None,
bucket_name: Optional[str] = None,
**kwargs,
):
# Translate request rate values into RequestRate objects
Expand Down Expand Up @@ -73,7 +73,7 @@ def __init__(
self.limit_statuses = limit_statuses
self.max_delay = max_delay
self.per_host = per_host
self.bucket = bucket
self.bucket_name = bucket_name
self._default_bucket = str(uuid4())

# If the superclass is an adapter or custom Session, pass along any valid keyword arguments
Expand All @@ -87,9 +87,8 @@ def send(self, request: PreparedRequest, **kwargs) -> Response:
Raises:
:py:exc:`.BucketFullException` if this request would result in a delay longer than ``max_delay``
"""
bucket_name = self.bucket or self._bucket_name(request)
with self.limiter.ratelimit(
bucket_name,
self._bucket_name(request),
delay=True,
max_delay=self.max_delay,
):
Expand Down
13 changes: 13 additions & 0 deletions test/test_requests_ratelimiter.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,19 @@ def test_custom_session(mock_sleep):
session.get(MOCKED_URL)
assert mock_sleep.called is True

@patch_sleep
def test_custom_bucket(mock_sleep):
"""With custom buckets, each session can be called independently without triggering rate limiting"""
session_a = get_mock_session(per_second=5, bucket_name="a")
session_b = get_mock_session(per_second=5, bucket_name="b")

for _ in range(5):
session_a.get(MOCKED_URL)
session_b.get(MOCKED_URL)
assert mock_sleep.called is False

session_a.get(MOCKED_URL)
assert mock_sleep.called is True

@patch_sleep
def test_429(mock_sleep):
Expand Down

0 comments on commit cf6d404

Please sign in to comment.