Avoid raising RepeatedConnectionLostError when listening frames (#52)

Author: vrslev

stompman/connection_manager.py

@@ -120,17 +120,13 @@ async def write_frame_reconnecting(self, frame: AnyClientFrame) -> None:
         raise RepeatedConnectionLostError(retry_attempts=self.connect_retry_attempts)
 
     async def read_frames_reconnecting(self) -> AsyncGenerator[AnyServerFrame, None]:
-        for _ in range(self.connect_retry_attempts):
+        while True:
             connection_state = await self._get_active_connection_state()
             try:
                 async for frame in connection_state.connection.read_frames():
                     yield frame
             except ConnectionLostError:
                 self._clear_active_connection_state()
-            else:
-                return
-
-        raise RepeatedConnectionLostError(retry_attempts=self.connect_retry_attempts)
 
     async def maybe_write_frame(self, frame: AnyClientFrame) -> bool:
         if not self._active_connection_state:

tests/test_connection_manager.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, AsyncIterable
 from types import SimpleNamespace
 from typing import Self
 from unittest import mock
@@ -235,20 +235,6 @@ class MockConnection(BaseMockConnection):
         await manager.write_frame_reconnecting(build_dataclass(ConnectFrame))
 
 
-async def test_read_frames_reconnecting_raises() -> None:
-    class MockConnection(BaseMockConnection):
-        @staticmethod
-        async def read_frames() -> AsyncGenerator[AnyServerFrame, None]:
-            raise ConnectionLostError
-            yield
-            await asyncio.sleep(0)
-
-    manager = EnrichedConnectionManager(connection_class=MockConnection)
-
-    with pytest.raises(RepeatedConnectionLostError):
-        [_ async for _ in manager.read_frames_reconnecting()]
-
-
 SIDE_EFFECTS = [(None,), (ConnectionLostError(), None), (ConnectionLostError(), ConnectionLostError(), None)]
 
 
@@ -299,11 +285,15 @@ async def read_frames() -> AsyncGenerator[AnyServerFrame, None]:
                 raise ConnectionLostError
             for frame in frames:
                 yield frame
-            await asyncio.sleep(0)
 
     manager = EnrichedConnectionManager(connection_class=MockConnection)
 
-    assert frames == [frame async for frame in manager.read_frames_reconnecting()]
+    async def take_all_frames() -> AsyncIterable[AnyServerFrame]:
+        iterator = manager.read_frames_reconnecting()
+        for _ in frames:
+            yield await anext(iterator)
+
+    assert frames == [frame async for frame in take_all_frames()]
 
 
 async def test_maybe_write_frame_connection_already_lost() -> None: