test: replace sleep with assert_eventually for rate reset verification

Author: CahidArda

tests/asyncio/test_flow_control.py

@@ -6,6 +6,7 @@
 from qstash import AsyncQStash
 from qstash.flow_control_api import GlobalParallelismInfo
 from qstash.message import FlowControl, PublishResponse
+from tests import assert_eventually_async
 
 
 @pytest.mark.asyncio
@@ -191,12 +192,13 @@ async def test_flow_control_reset_rate_async(async_client: AsyncQStash) -> None:
     # Reset the rate
     await async_client.flow_control.reset_rate(flow_control_key)
 
-    # sleep 1 second
-    await asyncio.sleep(1)
+    # We use assert_eventually because reset_rate doesn't take effect immediately.
+    # Not ideal but it's what we have.
+    async def check_rate_reset() -> None:
+        info = await async_client.flow_control.get(flow_control_key)
+        assert info.rate_count == 0
 
-    # Verify rate was reset by checking the flow control info
-    info = await async_client.flow_control.get(flow_control_key)
-    assert info.rate_count == 0
+    await assert_eventually_async(check_rate_reset)
 
     # Clean up
     await async_client.message.cancel(result.message_id)

tests/test_flow_control.py

@@ -3,6 +3,7 @@
 from qstash import QStash
 from qstash.flow_control_api import GlobalParallelismInfo
 from qstash.message import FlowControl, PublishResponse
+from tests import assert_eventually
 
 
 def test_flow_control_get(client: QStash) -> None:
@@ -180,12 +181,13 @@ def test_flow_control_reset_rate(client: QStash) -> None:
     # Reset the rate
     client.flow_control.reset_rate(flow_control_key)
 
-    # sleep 1 second
-    time.sleep(1)
+    # We use assert_eventually because reset_rate doesn't take effect immediately.
+    # Not ideal but it's what we have.
+    def check_rate_reset() -> None:
+        info = client.flow_control.get(flow_control_key)
+        assert info.rate_count == 0
 
-    # Verify rate was reset by checking the flow control info
-    info = client.flow_control.get(flow_control_key)
-    assert info.rate_count == 0
+    assert_eventually(check_rate_reset)
 
     # Clean up
     client.message.cancel(result.message_id)