diff --git a/src/httpx2/httpx2/_sse.py b/src/httpx2/httpx2/_sse.py index f04c4703..f420a22a 100644 --- a/src/httpx2/httpx2/_sse.py +++ b/src/httpx2/httpx2/_sse.py @@ -10,7 +10,7 @@ from collections.abc import AsyncIterator, Iterator from dataclasses import dataclass -from ._exceptions import TransportError +from ._exceptions import TransportError, request_context from ._models import Response __all__ = ["EventSource", "SSEError", "ServerSentEvent"] @@ -124,35 +124,34 @@ def response(self) -> Response: def _check_content_type(self) -> None: content_type, _, _ = self._response.headers.get("content-type", "").partition(";") if content_type.strip().lower() != "text/event-stream": - raise SSEError( - f"Expected response with content type 'text/event-stream', got {content_type.strip()!r}.", - request=self._response.request, - ) + raise SSEError(f"Expected response with content type 'text/event-stream', got {content_type.strip()!r}.") def __iter__(self) -> Iterator[ServerSentEvent]: - self._check_content_type() - decoder = _SSEDecoder() - lines = _SSELineDecoder() - for chunk in self._response.iter_text(): - for line in lines.decode(chunk): + with request_context(request=self._response._request): + self._check_content_type() + decoder = _SSEDecoder() + lines = _SSELineDecoder() + for chunk in self._response.iter_text(): + for line in lines.decode(chunk): + sse = decoder.decode(line) + if sse is not None: + yield sse + for line in lines.flush(): sse = decoder.decode(line) if sse is not None: yield sse - for line in lines.flush(): - sse = decoder.decode(line) - if sse is not None: - yield sse async def __aiter__(self) -> AsyncIterator[ServerSentEvent]: - self._check_content_type() - decoder = _SSEDecoder() - lines = _SSELineDecoder() - async for chunk in self._response.aiter_text(): - for line in lines.decode(chunk): + with request_context(request=self._response._request): + self._check_content_type() + decoder = _SSEDecoder() + lines = _SSELineDecoder() + async for chunk in self._response.aiter_text(): + for line in lines.decode(chunk): + sse = decoder.decode(line) + if sse is not None: + yield sse + for line in lines.flush(): sse = decoder.decode(line) if sse is not None: yield sse - for line in lines.flush(): - sse = decoder.decode(line) - if sse is not None: - yield sse diff --git a/tests/httpx2/test_sse.py b/tests/httpx2/test_sse.py index ae32fe2f..51a11c42 100644 --- a/tests/httpx2/test_sse.py +++ b/tests/httpx2/test_sse.py @@ -153,9 +153,11 @@ def handler(request: httpx2.Request) -> httpx2.Response: with httpx2.Client(transport=httpx2.MockTransport(handler)) as client: with client.sse("http://testserver/sse") as source: - with pytest.raises(httpx2.SSEError, match="text/event-stream"): + with pytest.raises(httpx2.SSEError, match="text/event-stream") as exc_info: list(source) + assert exc_info.value.request.url == "http://testserver/sse" + @pytest.mark.anyio async def test_content_type_mismatch_raises_async() -> None: @@ -164,9 +166,19 @@ def handler(request: httpx2.Request) -> httpx2.Response: async with httpx2.AsyncClient(transport=httpx2.MockTransport(handler)) as client: async with client.sse("http://testserver/sse") as source: - with pytest.raises(httpx2.SSEError, match="text/event-stream"): + with pytest.raises(httpx2.SSEError, match="text/event-stream") as exc_info: [event async for event in source] + assert exc_info.value.request.url == "http://testserver/sse" + + +def test_event_source_without_request() -> None: + response = httpx2.Response(200, content=b"data: hi\n\n", headers={"Content-Type": "text/event-stream"}) + + (event,) = list(httpx2.EventSource(response)) + + assert event.data == "hi" + def test_sets_sse_headers() -> None: captured: dict[str, str] = {}