From cf6d40488e4cc7835589817ed5fd175ae215cb21 Mon Sep 17 00:00:00 2001
From: Aaron Fulton <aaron@webtolife.software>
Date: Thu, 29 Feb 2024 19:30:01 +1300
Subject: [PATCH] Add unit test for bucket_name and fix associated issues

---
 README.md                                    |  3 +++
 requests_ratelimiter/requests_ratelimiter.py |  7 +++----
 test/test_requests_ratelimiter.py            | 13 +++++++++++++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 4c3315f..375897e 100644
--- a/README.md
+++ b/README.md
@@ -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')
diff --git a/requests_ratelimiter/requests_ratelimiter.py b/requests_ratelimiter/requests_ratelimiter.py
index 91473ea..e07921c 100644
--- a/requests_ratelimiter/requests_ratelimiter.py
+++ b/requests_ratelimiter/requests_ratelimiter.py
@@ -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
@@ -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
@@ -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,
         ):
diff --git a/test/test_requests_ratelimiter.py b/test/test_requests_ratelimiter.py
index 2ebfac4..388eed8 100644
--- a/test/test_requests_ratelimiter.py
+++ b/test/test_requests_ratelimiter.py
@@ -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):