86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import unittest
|
|
|
|
from backend.http.handler import HttpTransportMixin
|
|
|
|
|
|
class StreamError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class TransportStub(HttpTransportMixin):
|
|
def __init__(self) -> None:
|
|
self.statuses: list[int] = []
|
|
self.response_headers: list[tuple[str, str]] = []
|
|
self.wfile = io.BytesIO()
|
|
self.close_connection = False
|
|
|
|
def send_response(self, status: int) -> None:
|
|
self.statuses.append(int(status))
|
|
|
|
def send_header(self, name: str, value: str) -> None:
|
|
self.response_headers.append((name, value))
|
|
|
|
def end_headers(self) -> None:
|
|
pass
|
|
|
|
|
|
class HttpStreamingTests(unittest.TestCase):
|
|
def test_stream_transport_preserves_headers_events_and_completion(self) -> None:
|
|
handler = TransportStub()
|
|
handler.send_ndjson_stream(
|
|
({"type": "delta", "content": value} for value in ("甲", "乙")),
|
|
(StreamError,),
|
|
)
|
|
|
|
self.assertEqual(handler.statuses, [200])
|
|
self.assertEqual(
|
|
dict(handler.response_headers),
|
|
{
|
|
"Content-Type": "application/x-ndjson; charset=utf-8",
|
|
"Cache-Control": "no-cache, no-transform",
|
|
"X-Accel-Buffering": "no",
|
|
"Connection": "close",
|
|
},
|
|
)
|
|
events = [
|
|
json.loads(line)
|
|
for line in handler.wfile.getvalue().decode("utf-8").splitlines()
|
|
]
|
|
self.assertEqual(
|
|
events,
|
|
[
|
|
{"type": "delta", "content": "甲"},
|
|
{"type": "delta", "content": "乙"},
|
|
{"type": "done"},
|
|
],
|
|
)
|
|
self.assertTrue(handler.close_connection)
|
|
|
|
def test_stream_transport_preserves_feature_error_event(self) -> None:
|
|
def events():
|
|
yield {"type": "delta", "content": "partial"}
|
|
raise StreamError("stream failed")
|
|
|
|
handler = TransportStub()
|
|
handler.send_ndjson_stream(events(), (StreamError,))
|
|
payloads = [
|
|
json.loads(line)
|
|
for line in handler.wfile.getvalue().decode("utf-8").splitlines()
|
|
]
|
|
self.assertEqual(
|
|
payloads,
|
|
[
|
|
{"type": "delta", "content": "partial"},
|
|
{"type": "error", "error": "stream failed"},
|
|
],
|
|
)
|
|
self.assertTrue(handler.close_connection)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|