fix events-sent ping timeout (#69)

baize/asgi/responses.py

@@ -240,7 +240,7 @@ def __init__(
         self.charset = charset
 
     async def render_stream(self) -> AsyncGenerator[bytes, None]:
-        q: "asyncio.Queue[ServerSentEvent]" = asyncio.Queue(maxsize=1)
+        q: "asyncio.Queue[ServerSentEvent | None]" = asyncio.Queue(maxsize=1)
 
         should_stop = False
 
@@ -256,6 +256,7 @@ async def push() -> None:
                     except StopAsyncIteration:
                         should_stop = True
             finally:
+                await q.put(None)
                 g = self.iterable
                 if hasattr(g, "aclose"):
                     await g.aclose()  # type: ignore
@@ -266,6 +267,8 @@ async def push() -> None:
             while not (push_future.done() and q.empty()):
                 try:
                     event = await asyncio.wait_for(q.get(), timeout=self.ping_interval)
+                    if event is None:
+                        break
                     yield build_bytes_from_sse(event, self.charset)
                 except asyncio.TimeoutError:
                     yield b": ping\n\n"

baize/wsgi/responses.py

@@ -238,7 +238,7 @@ def __init__(
         self.charset = charset
 
     def render_stream(self) -> Generator[bytes, None, None]:
-        q: "queue.Queue[ServerSentEvent]" = queue.Queue(maxsize=1)
+        q: "queue.Queue[ServerSentEvent | None]" = queue.Queue(maxsize=1)
         should_stop = False
 
         def push() -> None:
@@ -253,6 +253,7 @@ def push() -> None:
                     except StopIteration:
                         should_stop = True
             finally:
+                q.put(None)
                 g = self.iterable
                 if hasattr(g, "close"):
                     g.close()  # type: ignore
@@ -263,6 +264,8 @@ def push() -> None:
             while not (push_future.done() and q.empty()):
                 try:
                     event = q.get(timeout=self.ping_interval)
+                    if event is None:
+                        break
                     yield build_bytes_from_sse(event, self.charset)
                 except queue.Empty:
                     yield b": ping\n\n"

tests/test_asgi.py

@@ -802,6 +802,7 @@ async def send_events() -> AsyncGenerator[ServerSentEvent, None]:
     ) as client:
         async with client.stream("GET", "/") as resp:
             resp.raise_for_status()
+            await resp.aread()
 
     assert killed
 

tests/test_wsgi.py

@@ -521,6 +521,7 @@ def send_events() -> Generator[ServerSentEvent, None, None]:
     ) as client:
         with client.stream("GET", "/") as resp:
             resp.raise_for_status()
+            resp.read()
 
     assert killed