test(workflow): make rate-limit burst and refill checks deterministic (#7433)

Author: masenf

tests/units/reflex_workflow/engine/test_runner.py

@@ -2005,27 +2005,36 @@ async def test_a_worker_with_one_slot_still_visits_every_table(session_factory):
         await engine.dispose()
 
 
-async def test_a_rate_limit_lets_through_a_burst_then_refills(session_factory):
+@pytest.mark.parametrize("enqueue_delay", [0, 1.1])
+async def test_a_rate_limit_lets_through_a_burst_then_refills(
+    session_factory, monkeypatch, enqueue_delay
+):
+    """A full bucket allows a burst and a partial refill allows two more runs."""
     STARTS.clear()
     provider = f"provider-{uuid.uuid4().hex}"
-    started = time.monotonic()
+    # Four tokens per hour prevent setup or scheduling delays from refilling
+    # the bucket; backdating it by half an hour releases exactly two tokens.
+    period = datetime.timedelta(hours=1)
+    monkeypatch.setattr(
+        Metered, "__workflow_limit__", Limit(by="provider", rate=4, per=period)
+    )
+    await asyncio.sleep(enqueue_delay)
     await start_many(Metered, provider, 8, provider=provider)
 
-    async def all_done() -> bool:
-        """Tell whether every run has finished.
-
-        Returns:
-            Whether it has.
-        """
+    for expected in (4, 6, 8):
+        if expected > 4:
+            async with session_factory() as session, session.begin():
+                await session.execute(
+                    update(WorkflowRate)
+                    .where(WorkflowRate.key == claim.bucket_key(Metered, provider))
+                    .values(tokens=0, updated_at=func.now() - period / 2)
+                )
+        # A settled worker has finished its steps and found no more to claim.
+        assert await runner.wake(datetime.timedelta(seconds=30))
+        assert len(STARTS[provider]) == expected
         rows = await Metered.by(Metered.provider == provider).all()
-        return len(rows) == 8 and all(row.status == "done" for row in rows)
-
-    await wait_until(all_done, timeout=60)
-    times = sorted(t - started for t in STARTS[provider])
-    # Four at once from a full bucket, then the rest as it refills: two a second.
-    assert len(times) == 8
-    assert times[3] < 1
-    assert times[7] >= 1.5
+        assert len(rows) == 8
+        assert sum(row.status == "done" for row in rows) == expected
 
 
 async def test_a_rate_limit_holds_across_workers(session_factory):