fix: ratelimiter ut (#710)

* `RateLimiter` 重构 * 修复在 `_query_per_period` 小于 1 时导致的 `RateLimiter` 不可用问题 --------- Co-authored-by: NuODaniel <[email protected]>
baidubce · Aug 3, 2024 · 4d12fe9 · 4d12fe9
1 parent 04fe59d
commit 4d12fe9
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/python/qianfan/resources/rate_limiter.py b/python/qianfan/resources/rate_limiter.py
@@ -446,13 +446,23 @@ def reset(
         query_per_period: float = 1,
         period_in_second: float = 1,
     ) -> None:
+        # 向上取整到 1，避免 SyncLimiter 失效
+        if query_per_period < 1:
+            period_in_second = period_in_second / query_per_period
+            query_per_period = 1
+
         self._sync_limiter.reset(query_per_period, period_in_second)
 
     async def async_reset(
         self,
         query_per_period: float = 1,
         period_in_second: float = 1,
     ) -> None:
+        # 向上取整到 1，避免 AsyncLimiter 失效
+        if query_per_period < 1:
+            period_in_second = period_in_second / query_per_period
+            query_per_period = 1
+
         self._async_limiter.max_rate = query_per_period
         self._async_limiter.time_period = period_in_second
         self._async_limiter._rate_per_sec = query_per_period / period_in_second

diff --git a/python/qianfan/tests/rate_limiter_test.py b/python/qianfan/tests/rate_limiter_test.py
@@ -228,7 +228,7 @@ def test_reset_once():
     rpm_rl = VersatileRateLimiter(query_per_second=5)
 
     assert not rpm_rl._is_rpm
-    assert rpm_rl._internal_qps_rate_limiter._query_per_period == 4.5
+    assert rpm_rl._internal_qps_rate_limiter._sync_limiter._query_per_period == 4.5
 
     def _reset_once():
         rpm_rl.reset_once(200)
@@ -245,15 +245,15 @@ def _reset_once():
     assert rpm_rl._has_been_reset
     assert not rpm_rl._is_rpm
     assert rpm_rl._new_query_per_second == 200 / 60
-    assert rpm_rl._internal_qps_rate_limiter._query_per_period == 3
+    assert rpm_rl._internal_qps_rate_limiter._sync_limiter._query_per_period == 3
 
 
 @pytest.mark.asyncio
 async def test_reset_once_async():
     rpm_rl = VersatileRateLimiter(request_per_minute=300)
 
     assert rpm_rl._is_rpm
-    assert rpm_rl._internal_rpm_rate_limiter._query_per_period == 270
+    assert rpm_rl._internal_rpm_rate_limiter._async_limiter.max_rate == 270
 
     awaitable_list = []
     for i in range(5):
@@ -264,7 +264,7 @@ async def test_reset_once_async():
     assert rpm_rl._has_been_reset
     assert rpm_rl._is_rpm
     assert rpm_rl._new_request_per_minute == 200
-    assert rpm_rl._internal_rpm_rate_limiter._query_per_period == 180
+    assert rpm_rl._internal_rpm_rate_limiter._async_limiter.max_rate == 180
 
 
 def test_reset_once_from_closed():
@@ -287,7 +287,7 @@ def _reset_once():
     assert rpm_rl._has_been_reset
     assert rpm_rl._is_rpm
     assert rpm_rl._new_request_per_minute == 200
-    assert rpm_rl._internal_rpm_rate_limiter._query_per_period == 180
+    assert rpm_rl._internal_rpm_rate_limiter._sync_limiter._query_per_period == 180
 
 
 @pytest.mark.asyncio
@@ -305,4 +305,4 @@ async def test_reset_once_async_from_closed():
     assert rpm_rl._has_been_reset
     assert rpm_rl._is_rpm
     assert rpm_rl._new_request_per_minute == 200
-    assert rpm_rl._internal_rpm_rate_limiter._query_per_period == 180
+    assert rpm_rl._internal_rpm_rate_limiter._sync_limiter._query_per_period == 180