Break out rate limiter protocol for testing

Author: blast-hardcheese

src/replit_river/rate_limiter.py

@@ -2,6 +2,7 @@
 import logging
 import random
 from contextvars import Context
+from typing import Protocol
 
 from replit_river.error_schema import RiverException
 from replit_river.transport_options import ConnectionRetryOptions
@@ -15,6 +16,13 @@ def __init__(self, code: str, message: str, client_id: str) -> None:
         self.client_id = client_id
 
 
+class RateLimiter(Protocol):
+    def start_restoring_budget(self, user: str) -> None: ...
+    def get_backoff_ms(self, user: str) -> float: ...
+    def has_budget(self, user: str) -> bool: ...
+    def consume_budget(self, user: str) -> None: ...
+
+
 class LeakyBucketRateLimit:
     """Asynchronous leaky bucket rate limiter.
 

src/replit_river/v2/session.py

@@ -54,7 +54,7 @@
     parse_transport_msg,
     send_transport_message,
 )
-from replit_river.rate_limiter import LeakyBucketRateLimit
+from replit_river.rate_limiter import RateLimiter
 from replit_river.rpc import (
     ACK_BIT,
     STREAM_OPEN_BIT,
@@ -143,7 +143,7 @@ class Session[HandshakeMetadata]:
     _wait_for_connected: asyncio.Event
 
     _client_id: str
-    _rate_limiter: LeakyBucketRateLimit
+    _rate_limiter: RateLimiter
     _uri_and_metadata_factory: Callable[
         [], Awaitable[UriAndMetadata[HandshakeMetadata]]
     ]
@@ -176,7 +176,7 @@ def __init__(
         transport_options: TransportOptions,
         close_session_callback: CloseSessionCallback,
         client_id: str,
-        rate_limiter: LeakyBucketRateLimit,
+        rate_limiter: RateLimiter,
         uri_and_metadata_factory: Callable[
             [], Awaitable[UriAndMetadata[HandshakeMetadata]]
         ],
@@ -1034,7 +1034,7 @@ async def _do_ensure_connected[HandshakeMetadata](
     client_id: str,
     session_id: str,
     server_id: str,
-    rate_limiter: LeakyBucketRateLimit,
+    rate_limiter: RateLimiter,
     uri_and_metadata_factory: Callable[
         [], Awaitable[UriAndMetadata[HandshakeMetadata]]
     ],