mock timer in rate limit test to make them deterministic

Author: arvidn

chia/_tests/core/server/test_rate_limits.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import asyncio
-
 import pytest
 from chia_rs.sized_ints import uint32
 
@@ -25,6 +23,28 @@
 test_different_versions_results: list[int] = []
 
 
+def sleep() -> float:
+    return 0.0
+
+
+@pytest.fixture
+def mock_timer(monkeypatch):
+    current_time = 1000.0
+
+    def mock_monotonic(*args, **kwargs) -> float:
+        nonlocal current_time
+        return current_time
+
+    def mock_sleep(*args, **kwargs) -> None:
+        nonlocal current_time
+        current_time += int(args[0])
+
+    import chia
+
+    monkeypatch.setattr(chia.server.rate_limits, "get_time", mock_monotonic)
+    monkeypatch.setattr(chia._tests.core.server.test_rate_limits, "sleep", mock_sleep)
+
+
 class TestRateLimits:
     @pytest.mark.anyio
     async def test_get_rate_limits_to_use(self):
@@ -33,7 +53,7 @@ async def test_get_rate_limits_to_use(self):
         assert get_rate_limits_to_use(rl_v1, rl_v1) == get_rate_limits_to_use(rl_v1, rl_v2)
 
     @pytest.mark.anyio
-    async def test_too_many_messages(self):
+    async def test_too_many_messages(self, mock_timer):
         # Too many messages
         r = RateLimiter(incoming=True)
         new_tx_message = make_msg(ProtocolMessageTypes.new_transaction, bytes([1] * 40))
@@ -61,7 +81,7 @@ async def test_too_many_messages(self):
         assert saw_disconnect
 
     @pytest.mark.anyio
-    async def test_large_message(self):
+    async def test_large_message(self, mock_timer):
         # Large tx
         small_tx_message = make_msg(ProtocolMessageTypes.respond_transaction, bytes([1] * 500 * 1024))
         large_tx_message = make_msg(ProtocolMessageTypes.new_transaction, bytes([1] * 3 * 1024 * 1024))
@@ -81,7 +101,7 @@ async def test_large_message(self):
         assert r.process_msg_and_check(large_blocks_message, rl_v2, rl_v2) is not None
 
     @pytest.mark.anyio
-    async def test_too_much_data(self):
+    async def test_too_much_data(self, mock_timer):
         # Too much data
         r = RateLimiter(incoming=True)
         tx_message = make_msg(ProtocolMessageTypes.respond_transaction, bytes([1] * 500 * 1024))
@@ -108,7 +128,7 @@ async def test_too_much_data(self):
         assert saw_disconnect
 
     @pytest.mark.anyio
-    async def test_non_tx_aggregate_limits(self):
+    async def test_non_tx_aggregate_limits(self, mock_timer):
         # Frequency limits
         r = RateLimiter(incoming=True)
         message_1 = make_msg(ProtocolMessageTypes.coin_state_update, bytes([1] * 32))
@@ -144,7 +164,7 @@ async def test_non_tx_aggregate_limits(self):
         assert saw_disconnect
 
     @pytest.mark.anyio
-    async def test_periodic_reset(self):
+    async def test_periodic_reset(self, mock_timer):
         r = RateLimiter(True, 5)
         tx_message = make_msg(ProtocolMessageTypes.respond_transaction, bytes([1] * 500 * 1024))
         for i in range(10):
@@ -157,7 +177,7 @@ async def test_periodic_reset(self):
                 saw_disconnect = True
         assert saw_disconnect
         assert r.process_msg_and_check(tx_message, rl_v2, rl_v2) is not None
-        await asyncio.sleep(6)
+        sleep(6)
         assert r.process_msg_and_check(tx_message, rl_v2, rl_v2) is None
 
         # Counts reset also
@@ -172,11 +192,11 @@ async def test_periodic_reset(self):
             if response is not None:
                 saw_disconnect = True
         assert saw_disconnect
-        await asyncio.sleep(6)
+        sleep(6)
         assert r.process_msg_and_check(new_tx_message, rl_v2, rl_v2) is None
 
     @pytest.mark.anyio
-    async def test_percentage_limits(self):
+    async def test_percentage_limits(self, mock_timer):
         r = RateLimiter(True, 60, 40)
         new_peak_message = make_msg(ProtocolMessageTypes.new_peak, bytes([1] * 40))
         for i in range(50):
@@ -235,7 +255,7 @@ async def test_percentage_limits(self):
         assert saw_disconnect
 
     @pytest.mark.anyio
-    async def test_too_many_outgoing_messages(self):
+    async def test_too_many_outgoing_messages(self, mock_timer):
         # Too many messages
         r = RateLimiter(incoming=False)
         new_peers_message = make_msg(ProtocolMessageTypes.respond_peers, bytes([1]))
@@ -258,7 +278,7 @@ async def test_too_many_outgoing_messages(self):
         assert r.process_msg_and_check(new_signatures_message, rl_v2, rl_v2) is None
 
     @pytest.mark.anyio
-    async def test_too_many_incoming_messages(self):
+    async def test_too_many_incoming_messages(self, mock_timer):
         # Too many messages
         r = RateLimiter(incoming=True)
         new_peers_message = make_msg(ProtocolMessageTypes.respond_peers, bytes([1]))
@@ -353,7 +373,7 @@ async def test_different_versions(self, node_with_params, node_with_params_b, se
             assert len(set(test_different_versions_results)) >= 2
 
     @pytest.mark.anyio
-    async def test_compose(self):
+    async def test_compose(self, mock_timer):
         rl_1 = rl_numbers[1]
         rl_2 = rl_numbers[2]
         assert ProtocolMessageTypes.respond_children in rl_1["rate_limits_other"]

chia/server/rate_limits.py

@@ -14,6 +14,10 @@
 log = logging.getLogger(__name__)
 
 
+def get_time() -> float:
+    return monotonic()
+
+
 # TODO: only full node disconnects based on rate limits
 class RateLimiter:
     incoming: bool
@@ -35,7 +39,7 @@ def __init__(self, incoming: bool, reset_seconds: int = 60, percentage_of_limit:
         """
         self.incoming = incoming
         self.reset_seconds = reset_seconds
-        self.current_slot = int(monotonic() // reset_seconds)
+        self.current_slot = int(get_time() // reset_seconds)
         self.message_counts = Counter()
         self.message_cumulative_sizes = Counter()
         self.percentage_of_limit = percentage_of_limit
@@ -51,7 +55,7 @@ def process_msg_and_check(
         hit and the message is good to be sent or received.
         """
 
-        current_slot = int(monotonic() // self.reset_seconds)
+        current_slot = int(get_time() // self.reset_seconds)
         if current_slot != self.current_slot:
             self.current_slot = current_slot
             self.message_counts = Counter()