1616"""Tests for async stream."""
1717
1818import asyncio
19- from typing import List
19+ import io
20+ import json
21+ from typing import Any , List
2022from unittest import mock
2123from unittest .mock import AsyncMock
2224from unittest .mock import MagicMock
2325from unittest .mock import patch
2426import pytest
27+ import requests
2528
2629
2730try :
@@ -51,15 +54,12 @@ class MockHTTPXResponse(httpx.Response):
5154 """Mock httpx.Response class for testing."""
5255
5356 def __init__ (self , lines : List [str ]):
54- self .aiter_lines = MagicMock ()
55- self .aiter_lines .return_value .__aiter__ = MagicMock (
56- return_value = self ._async_line_iterator (lines )
57- )
57+ self .aiter_bytes = MagicMock (return_value = self ._async_byte_iterator (lines ))
5858 self .aclose = AsyncMock ()
5959
60- async def _async_line_iterator (self , lines : List [str ]):
60+ async def _async_byte_iterator (self , lines : List [str ]):
6161 for line in lines :
62- yield line
62+ yield line . encode ( 'utf-8' ) + b' \n '
6363
6464
6565class MockAIOHTTPResponse (aiohttp .ClientResponse ):
@@ -128,7 +128,7 @@ async def test_httpx_simple_lines(responses: api_client.HttpResponse):
128128 results = [line async for line in responses ._aiter_response_stream ()]
129129
130130 assert results == lines
131- mock_response .aiter_lines .assert_called_once ()
131+ mock_response .aiter_bytes .assert_called_once ()
132132 mock_response .aclose .assert_called_once ()
133133
134134
@@ -141,7 +141,7 @@ async def test_httpx_data_prefix(responses: api_client.HttpResponse):
141141 results = [line async for line in responses ._aiter_response_stream ()]
142142
143143 assert results == ["{ 'message': 'hello' }" , "{ 'status': 'ok' }" ]
144- mock_response .aiter_lines .assert_called_once ()
144+ mock_response .aiter_bytes .assert_called_once ()
145145 mock_response .aclose .assert_called_once ()
146146
147147@pytest .mark .asyncio
@@ -158,7 +158,7 @@ async def test_httpx_multiline_data(responses: api_client.HttpResponse) -> None:
158158 results = [line async for line in responses ._aiter_response_stream ()]
159159
160160 assert results == ["{\n 'message': 'hello'\n }" ]
161- mock_response .aiter_lines .assert_called_once ()
161+ mock_response .aiter_bytes .assert_called_once ()
162162 mock_response .aclose .assert_called_once ()
163163
164164
@@ -170,7 +170,9 @@ def test_httpx_multiline_data_sync(responses: api_client.HttpResponse) -> None:
170170 "" ,
171171 ]
172172 mock_response = MagicMock (spec = httpx .Response )
173- mock_response .iter_lines .return_value = iter (lines )
173+ mock_response .iter_bytes .return_value = iter (
174+ [b'\n ' .join (line .encode ('utf-8' ) for line in lines ) + b'\n ' ]
175+ )
174176 responses .response_stream = mock_response
175177
176178 results = list (responses ._iter_response_stream ())
@@ -193,7 +195,7 @@ async def test_httpx_multiple_json_chunk(responses: api_client.HttpResponse):
193195 results = [line async for line in responses ._aiter_response_stream ()]
194196
195197 assert results == ['{ "id": 1 }' , '{ "id": 2 }' , '{ "id": 3 }' ]
196- mock_response .aiter_lines .assert_called_once ()
198+ mock_response .aiter_bytes .assert_called_once ()
197199 mock_response .aclose .assert_called_once ()
198200
199201
@@ -207,7 +209,7 @@ async def test_httpx_incomplete_json_at_end(responses: api_client.HttpResponse):
207209
208210 # The remaining chunk is yielded
209211 assert results == ['{ "partial": "data"' ]
210- mock_response .aiter_lines .assert_called_once ()
212+ mock_response .aiter_bytes .assert_called_once ()
211213 mock_response .aclose .assert_called_once ()
212214
213215
@@ -220,7 +222,7 @@ async def test_httpx_empty_stream(responses: api_client.HttpResponse):
220222 results = [line async for line in responses ._aiter_response_stream ()]
221223
222224 assert results == []
223- mock_response .aiter_lines .assert_called_once ()
225+ mock_response .aiter_bytes .assert_called_once ()
224226 mock_response .aclose .assert_called_once ()
225227
226228
@@ -538,3 +540,185 @@ async def test_error_event_in_generate_content_stream_async(mock_transport):
538540
539541 with pytest .raises (errors .ServerError ):
540542 await anext (generated_response )
543+
544+
545+ _LINE_SEPARATORS = pytest .mark .parametrize (
546+ 'sep' , ['\u2028 ' , '\u2029 ' , '\x85 ' ], ids = ['u2028' , 'u2029' , 'u0085' ]
547+ )
548+
549+
550+ def _sse_body_with_separator (sep : str ) -> bytes :
551+ return b'' .join (
552+ b'data: '
553+ + json .dumps (c , ensure_ascii = False ).encode ('utf-8' )
554+ + b'\r \n \r \n '
555+ for c in _sse_events_with_separator (sep )
556+ )
557+
558+
559+ def _sse_events_with_separator (sep : str ) -> list [dict [str , Any ]]:
560+ return [
561+ {
562+ 'candidates' : [{
563+ 'content' : {
564+ 'role' : 'model' ,
565+ 'parts' : [{'text' : f'line one{ sep } line two' }],
566+ }
567+ }]
568+ },
569+ {
570+ 'candidates' : [{
571+ 'content' : {'role' : 'model' , 'parts' : [{'text' : ' done' }]},
572+ 'finishReason' : 'STOP' ,
573+ }]
574+ },
575+ ]
576+
577+
578+ def _sse_handler (sep : str ):
579+ def handler (_ : httpx .Request ) -> httpx .Response :
580+ return httpx .Response (
581+ 200 ,
582+ content = _sse_body_with_separator (sep ),
583+ headers = {'content-type' : 'text/event-stream' },
584+ )
585+
586+ return handler
587+
588+
589+ @_LINE_SEPARATORS
590+ def test_generate_content_stream_line_separator_in_text (sep ):
591+ """Unicode line separators inside a JSON string do not split an SSE line."""
592+ client = Client (
593+ api_key = 'test_api_key' ,
594+ http_options = types .HttpOptions (
595+ client_args = {'transport' : httpx .MockTransport (_sse_handler (sep ))}
596+ ),
597+ )
598+
599+ stream = client .models .generate_content_stream (
600+ model = 'gemini-2.5-flash' , contents = 'hi'
601+ )
602+
603+ assert [c .text for c in stream ] == [f'line one{ sep } line two' , ' done' ]
604+
605+
606+ @_LINE_SEPARATORS
607+ @pytest .mark .asyncio
608+ async def test_generate_content_stream_line_separator_in_text_async (sep ):
609+ """Async httpx streams keep Unicode line separators inside JSON strings."""
610+ client = Client (
611+ api_key = 'test_api_key' ,
612+ http_options = types .HttpOptions (
613+ async_client_args = {
614+ 'transport' : httpx .MockTransport (_sse_handler (sep ))
615+ }
616+ ),
617+ )
618+
619+ stream = await client .aio .models .generate_content_stream (
620+ model = 'gemini-2.5-flash' , contents = 'hi'
621+ )
622+
623+ assert [c .text async for c in stream ] == [
624+ f'line one{ sep } line two' ,
625+ ' done' ,
626+ ]
627+
628+
629+ @_LINE_SEPARATORS
630+ def test_requests_stream_line_separator_in_text (sep ):
631+ """The requests transport splits SSE lines only at CR, LF and CRLF."""
632+ resp = requests .Response ()
633+ resp .status_code = 200
634+ resp .raw = io .BytesIO (_sse_body_with_separator (sep ))
635+ resp .encoding = 'utf-8'
636+
637+ results = list (
638+ api_client .HttpResponse (
639+ headers = {}, response_stream = resp
640+ )._iter_response_stream ()
641+ )
642+
643+ assert results == [
644+ json .dumps (c , ensure_ascii = False ) for c in _sse_events_with_separator (sep )
645+ ]
646+
647+
648+ def test_httpx_stream_crlf_split_across_chunks ():
649+ """A CRLF split across two chunks ends exactly one line."""
650+ chunks = [b'data: {"a": 1}\r ' , b'\n \r \n ' , b'data: {"b": 2}\r \n \r \n ' ]
651+ resp = api_client .HttpResponse (
652+ headers = {}, response_stream = mock_response (chunks )
653+ )
654+
655+ assert list (resp ._iter_response_stream ()) == ['{"a": 1}' , '{"b": 2}' ]
656+
657+
658+ def test_httpx_stream_cr_only_terminators ():
659+ """A bare CR ends a line, as the SSE format requires."""
660+ chunks = [b'data: {"a": 1}\r \r data: {"b": 2}\r \r ' ]
661+ resp = api_client .HttpResponse (
662+ headers = {}, response_stream = mock_response (chunks )
663+ )
664+
665+ assert list (resp ._iter_response_stream ()) == ['{"a": 1}' , '{"b": 2}' ]
666+
667+
668+ def test_httpx_stream_multibyte_char_split_across_chunks ():
669+ """A UTF-8 character split across chunks decodes as one character."""
670+ chunks = [b'data: {"t": "x\xe2 \x80 ' , b'\xa8 y"}\n \n ' ]
671+ resp = api_client .HttpResponse (
672+ headers = {}, response_stream = mock_response (chunks )
673+ )
674+
675+ assert list (resp ._iter_response_stream ()) == ['{"t": "x\u2028 y"}' ]
676+
677+
678+ def test_httpx_stream_multiline_data_with_line_separator ():
679+ """Multi-line data fields still join with LF around a U+2028."""
680+ chunks = [b'data: {"t": "a\xe2 \x80 \xa8 b",\n ' , b'data: "u": 1}\n \n ' ]
681+ resp = api_client .HttpResponse (
682+ headers = {}, response_stream = mock_response (chunks )
683+ )
684+
685+ assert list (resp ._iter_response_stream ()) == ['{"t": "a\u2028 b",\n "u": 1}' ]
686+
687+
688+ def test_httpx_stream_invalid_utf8_is_replaced ():
689+ """Invalid UTF-8 bytes become U+FFFD instead of raising."""
690+ chunks = [b'data: {"t": "\xff "}\n \n ' ]
691+ resp = api_client .HttpResponse (
692+ headers = {}, response_stream = mock_response (chunks )
693+ )
694+
695+ assert list (resp ._iter_response_stream ()) == ['{"t": "\ufffd "}' ]
696+
697+
698+ def test_httpx_stream_final_line_without_terminator ():
699+ """A last line with no terminator is still yielded at end of stream."""
700+ chunks = [b'data: {"a": 1}\n \n {"error": ' , b'{"code": 500}}' ]
701+ resp = api_client .HttpResponse (
702+ headers = {}, response_stream = mock_response (chunks )
703+ )
704+
705+ assert list (resp ._iter_response_stream ()) == [
706+ '{"a": 1}' ,
707+ '{"error": {"code": 500}}' ,
708+ ]
709+
710+
711+ @pytest .mark .asyncio
712+ async def test_httpx_async_stream_line_separator_in_text (
713+ responses : api_client .HttpResponse ,
714+ ):
715+ """The async httpx branch keeps a U+2028 line whole and closes once."""
716+ httpx_response = MockHTTPXResponse (
717+ ['data: {"t": "a\u2028 b"}' , '' , 'data: {"t": "c"}' , '' ]
718+ )
719+ responses .response_stream = httpx_response
720+
721+ results = [line async for line in responses ._aiter_response_stream ()]
722+
723+ assert results == ['{"t": "a\u2028 b"}' , '{"t": "c"}' ]
724+ httpx_response .aclose .assert_called_once ()
0 commit comments