Reraise ConnectionLostError on uvloop caused RuntimeError when writing heartbeat (#28)

Author: vrslev

.github/workflows/publish.yml

@@ -18,7 +18,7 @@ jobs:
           path: ~/.cache/uv
           key: publish-${{ hashFiles('pyproject.toml') }}
       - uses: extractions/setup-just@v2
-      - run: curl -LsSf https://astral.sh/uv/install.sh | sh
+      - run: curl -LsSf https://astral.sh/uv/0.2.11/install.sh | sh
       - run: just publish
         env:
           PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}

.github/workflows/test.yml

@@ -26,7 +26,7 @@ jobs:
         with:
           path: ~/.cache/uv
           key: check-types-${{ hashFiles('pyproject.toml') }}
-      - run: curl -LsSf https://astral.sh/uv/install.sh | sh
+      - run: curl -LsSf https://astral.sh/uv/0.2.11/install.sh | sh
       - uses: extractions/setup-just@v2
       - run: just install check-types
 
@@ -41,7 +41,7 @@ jobs:
         with:
           path: ~/.cache/uv
           key: lint-${{ hashFiles('pyproject.toml') }}
-      - run: curl -LsSf https://astral.sh/uv/install.sh | sh
+      - run: curl -LsSf https://astral.sh/uv/0.2.11/install.sh | sh
       - uses: extractions/setup-just@v2
       - run: just install lint
 
@@ -62,7 +62,7 @@ jobs:
         with:
           path: ~/.cache/uv
           key: ${{ github.job }}-${{ matrix.python-version }}-${{ hashFiles('pyproject.toml') }}
-      - run: curl -LsSf https://astral.sh/uv/install.sh | sh
+      - run: curl -LsSf https://astral.sh/uv/0.2.11/install.sh | sh
       - uses: extractions/setup-just@v2
       - run: just install test -vv
 

Dockerfile

@@ -2,7 +2,7 @@ ARG PYTHON_VERSION
 FROM python:${PYTHON_VERSION}-slim-bullseye
 
 # hadolint ignore=DL3013, DL3042
-RUN --mount=type=cache,target=~/.cache/pip pip install uv
+RUN --mount=type=cache,target=~/.cache/pip pip install uv==0.2.11
 
 WORKDIR /app
 

stompman/connection.py

@@ -56,7 +56,8 @@ async def close(self) -> None:
             await self.writer.wait_closed()
 
     def write_heartbeat(self) -> None:
-        return self.writer.write(NEWLINE)
+        with _reraise_connection_lost(RuntimeError):
+            return self.writer.write(NEWLINE)
 
     async def write_frame(self, frame: AnyClientFrame) -> None:
         with _reraise_connection_lost(RuntimeError):

tests/integration.py

@@ -1,5 +1,7 @@
 import asyncio
 import os
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
 from uuid import uuid4
 
 import pytest
@@ -49,7 +51,35 @@ async def consume() -> None:
         task_group.create_task(produce())
 
 
-async def test_raises_connection_lost_error(server: stompman.ConnectionParameters) -> None:
+@asynccontextmanager
+async def closed_client(server: stompman.ConnectionParameters) -> AsyncGenerator[stompman.Client, None]:
+    async with stompman.Client(servers=[server], read_timeout=10, connection_confirmation_timeout=10) as client:
+        await client._connection.close()
+        yield client
+
+
+async def test_raises_connection_lost_error_in_aexit(server: stompman.ConnectionParameters) -> None:
+    with pytest.raises(ConnectionLostError):
+        async with closed_client(server):
+            pass
+
+
+async def test_raises_connection_lost_error_in_write_frame(server: stompman.ConnectionParameters) -> None:
+    client = await closed_client(server).__aenter__()  # noqa: PLC2801
+
+    with pytest.raises(ConnectionLostError):
+        await client._connection.write_frame(stompman.ConnectFrame(headers={"accept-version": "", "host": ""}))
+
+    with pytest.raises(ConnectionLostError):
+        await client.__aexit__(None, None, None)
+
+
+@pytest.mark.parametrize("anyio_backend", [("asyncio", {"use_uvloop": True})])
+async def test_raises_connection_lost_error_in_write_heartbeat(server: stompman.ConnectionParameters) -> None:
+    client = await closed_client(server).__aenter__()  # noqa: PLC2801
+
+    with pytest.raises(ConnectionLostError):
+        client._connection.write_heartbeat()
+
     with pytest.raises(ConnectionLostError):
-        async with stompman.Client(servers=[server], read_timeout=10, connection_confirmation_timeout=10) as consumer:
-            await consumer._connection.close()
+        await client.__aexit__(None, None, None)

tests/test_connection.py

@@ -108,6 +108,15 @@ class MockWriter:
         await connection.close()
 
 
+async def test_connection_write_heartbeat_runtime_error(monkeypatch: pytest.MonkeyPatch) -> None:
+    class MockWriter:
+        write = mock.Mock(side_effect=RuntimeError)
+
+    connection = await make_mocked_connection(monkeypatch, mock.Mock(), MockWriter())
+    with pytest.raises(ConnectionLostError):
+        connection.write_heartbeat()
+
+
 async def test_connection_write_frame_connection_error(monkeypatch: pytest.MonkeyPatch) -> None:
     class MockWriter:
         write = mock.Mock()