tests/core/test_sse.py¶
Source from this local checkout, regenerated when the reader rebuilds.
Line links use #L<number>; a GitHub line range opens its first line.
1 # Copyright 2025 Softwell S.r.l.2 #3 # Licensed under the Apache License, Version 2.0 (the "License");4 # you may not use this file except in compliance with the License.5 # You may obtain a copy of the License at6 #7 # https://www.apache.org/licenses/LICENSE-2.08 #9 # Unless required by applicable law or agreed to in writing, software10 # distributed under the License is distributed on an "AS IS" BASIS,11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12 # See the License for the specific language governing permissions and13 # limitations under the License.14 15 """Tests for SseStream (core 1e Phase 5): text/event-stream framing.16 17 Real objects, no mocks: frame single events directly, iterate a real async18 source into wire bytes, exercise the heartbeat on a genuinely silent source19 (a small ``keepalive_seconds``, no clock faking), and confirm the20 ``StreamingResponse`` wrapper carries the SSE headers.21 """22 23 from __future__ import annotations24 25 import asyncio26 from collections.abc import AsyncIterable27 from typing import Any28 29 from genro_asgi.sse import KEEPALIVE_SECONDS, SseStream30 from genro_asgi.streaming import StreamingResponse31 from genro_asgi.types import Message, Scope32 33 34 async def source(*events: dict[str, Any]) -> AsyncIterable[dict[str, Any]]:35 """An async iterator over the given event dicts."""36 for event in events:37 yield event38 39 40 class TestFrame:41 """One event dict -> one SSE record (ends with a blank line)."""42 43 def test_data_only(self) -> None:44 frame = SseStream(source()).frame({"data": "hello"})45 assert frame == b"data: hello\n\n"46 47 def test_id_and_event(self) -> None:48 frame = SseStream(source()).frame({"id": "7", "event": "progress", "data": "x"})49 assert frame == b"id: 7\nevent: progress\ndata: x\n\n"50 51 def test_non_string_data_is_json(self) -> None:52 frame = SseStream(source()).frame({"data": {"pct": 50}})53 assert frame == b'data: {"pct": 50}\n\n'54 55 def test_multiline_data_split(self) -> None:56 frame = SseStream(source()).frame({"data": "line1\nline2"})57 assert frame == b"data: line1\ndata: line2\n\n"58 59 def test_id_omitted_when_absent(self) -> None:60 frame = SseStream(source()).frame({"event": "ping", "data": "1"})61 assert b"id:" not in frame62 63 64 class TestIteration:65 """Iterating the stream yields wire bytes for each event."""66 67 async def test_events_framed_in_order(self) -> None:68 stream = SseStream(source({"data": "a"}, {"data": "b"}))69 chunks = [chunk async for chunk in stream]70 assert chunks == [b"data: a\n\n", b"data: b\n\n"]71 72 async def test_retry_emitted_once_at_start(self) -> None:73 stream = SseStream(source({"data": "a"}), retry_ms=5000)74 chunks = [chunk async for chunk in stream]75 assert chunks[0] == b"retry: 5000\n\n"76 assert chunks[1] == b"data: a\n\n"77 assert sum(c.startswith(b"retry:") for c in chunks) == 178 79 async def test_exhausted_source_ends_stream(self) -> None:80 stream = SseStream(source())81 chunks = [chunk async for chunk in stream]82 assert chunks == []83 84 85 class TestHeartbeat:86 """A silent source past keepalive_seconds emits ``: keepalive`` comments."""87 88 async def test_keepalive_while_silent_then_event(self) -> None:89 async def slow() -> AsyncIterable[dict[str, Any]]:90 await asyncio.sleep(0.12) # silent longer than keepalive91 yield {"data": "late"}92 93 stream = SseStream(slow(), keepalive_seconds=0.04)94 chunks = [chunk async for chunk in stream]95 keepalives = [c for c in chunks if c == b": keepalive\n\n"]96 assert len(keepalives) >= 2 # at least two idle intervals97 assert chunks[-1] == b"data: late\n\n" # the real event still arrives98 99 async def test_no_keepalive_when_source_is_prompt(self) -> None:100 stream = SseStream(source({"data": "a"}), keepalive_seconds=1.0)101 chunks = [chunk async for chunk in stream]102 assert b": keepalive\n\n" not in chunks103 104 def test_default_keepalive_interval(self) -> None:105 assert SseStream(source()).keepalive_seconds == KEEPALIVE_SECONDS106 107 108 class TestConsumerGone:109 """Cancelling the consumer closes the source (its ``finally`` must run)."""110 111 async def test_source_finalized_on_cancel(self) -> None:112 closed = asyncio.Event()113 114 async def endless() -> AsyncIterable[dict[str, Any]]:115 try:116 yield {"data": "first"}117 await asyncio.Event().wait() # blocks forever118 yield {"data": "never"}119 finally:120 closed.set() # the unsubscribe seam121 122 async def consume() -> None:123 async for _ in SseStream(endless(), keepalive_seconds=60.0):124 task_started.set()125 126 task_started = asyncio.Event()127 task = asyncio.get_running_loop().create_task(consume())128 await task_started.wait() # first frame arrived129 task.cancel()130 try:131 await task132 except asyncio.CancelledError:133 pass134 assert closed.is_set() # finally ran, no leaked read135 136 137 class TestResponseWrapper:138 """``response()`` wraps the stream in a StreamingResponse with SSE headers."""139 140 async def test_sse_headers_set(self) -> None:141 stream = SseStream(source({"data": "a"}))142 response = stream.response()143 assert isinstance(response, StreamingResponse)144 145 scope: Scope = {"type": "http", "method": "GET", "path": "/", "headers": []}146 sent: list[Message] = []147 148 async def receive() -> Message:149 return {"type": "http.request"}150 151 async def send(message: Message) -> None:152 sent.append(message)153 154 await response(scope, receive, send)155 headers = dict(sent[0]["headers"])156 # text/* gets the charset appended, uniform with Response (SSE is UTF-8)157 assert headers[b"content-type"] == b"text/event-stream; charset=utf-8"158 assert headers[b"cache-control"] == b"no-cache"159 assert headers[b"connection"] == b"keep-alive"160 # the event and the terminal body come through the stream transport161 assert any(m.get("body") == b"data: a\n\n" for m in sent)162 assert sent[-1]["more_body"] is False