Don't swallow asyncio.CancelledError

aiohttp_sse/__init__.py

@@ -1,7 +1,7 @@
 import asyncio
 import io
 import re
-from contextlib import suppress
+import sys
 from types import TracebackType
 from typing import Any, Mapping, Optional, Type, TypeVar, Union, overload
 
@@ -142,8 +142,16 @@
         """
         if self._ping_task is None:
             raise RuntimeError("Response is not started")
-        with suppress(asyncio.CancelledError):
+
+        try:
             await self._ping_task
+        except asyncio.CancelledError:
+            if (
+                sys.version_info >= (3, 11)
+                and (task := asyncio.current_task())
+                and task.cancelling()
+            ):
+                raise
 
     def stop_streaming(self) -> None:
         """Used in conjunction with ``wait`` could be called from other task

tests/test_sse.py

@@ -1,4 +1,5 @@
 import asyncio
+import sys
 from typing import Awaitable, Callable, List
 
 import pytest
@@ -525,3 +526,36 @@
         streamed_data = await resp.text()
 
     assert streamed_data == "data: foo\r\n\r\n"
+
+
+@pytest.mark.skipif(
+    sys.version_info < (3, 11),
+    reason=".cancelling() missing in older versions",
+)
+async def test_cancelled_not_swallowed(aiohttp_client: ClientFixture) -> None:
+    """Test asyncio.CancelledError is not swallowed by .wait().
+
+    Relates to:
+    https://github.com/aio-libs/aiohttp-sse/issues/458
+    """
+
+    async def endless_task(sse: EventSourceResponse) -> None:
+        while True:
+            await sse.wait()
+
+    async def handler(request: web.Request) -> EventSourceResponse:
+        async with sse_response(request) as sse:
+            task = asyncio.create_task(endless_task(sse))
+            await asyncio.sleep(0)
+            task.cancel()
+            await task
+
+        return sse  # pragma: no cover
+
+    app = web.Application()
+    app.router.add_route("GET", "/", handler)
+
+    client = await aiohttp_client(app)
+
+    async with client.get("/") as response:
+        assert 200 == response.status