e6a02cb90d
* init * enrich doc * init * init * init * doc * chore: remove deprecated email field from pyproject.toml --------- Co-authored-by: zhengweijun <weijun.zheng@aminer.cn>
92 lines
1.8 KiB
Python
92 lines
1.8 KiB
Python
from typing import Iterator
|
|
|
|
import pytest
|
|
|
|
from zai.core._streaming import SSELineParser
|
|
|
|
|
|
def test_basic() -> None:
|
|
def body() -> Iterator[str]:
|
|
yield 'event: completion'
|
|
yield 'data: {"foo":true}'
|
|
yield ''
|
|
|
|
it = SSELineParser().iter_lines(body())
|
|
sse = next(it)
|
|
assert sse.event == 'completion'
|
|
assert sse.json_data() == {'foo': True}
|
|
|
|
with pytest.raises(StopIteration):
|
|
next(it)
|
|
|
|
|
|
def test_data_missing_event() -> None:
|
|
def body() -> Iterator[str]:
|
|
yield 'data: {"foo":true}'
|
|
yield ''
|
|
|
|
it = SSELineParser().iter_lines(body())
|
|
sse = next(it)
|
|
assert sse.event is None
|
|
assert sse.json_data() == {'foo': True}
|
|
|
|
with pytest.raises(StopIteration):
|
|
next(it)
|
|
|
|
|
|
def test_event_missing_data() -> None:
|
|
def body() -> Iterator[str]:
|
|
yield 'event: ping'
|
|
yield ''
|
|
|
|
it = SSELineParser().iter_lines(body())
|
|
sse = next(it)
|
|
assert sse.event == 'ping'
|
|
assert sse.data == ''
|
|
|
|
with pytest.raises(StopIteration):
|
|
next(it)
|
|
|
|
|
|
def test_multiple_events() -> None:
|
|
def body() -> Iterator[str]:
|
|
yield 'event: ping'
|
|
yield ''
|
|
yield 'event: completion'
|
|
yield ''
|
|
|
|
it = SSELineParser().iter_lines(body())
|
|
|
|
sse = next(it)
|
|
assert sse.event == 'ping'
|
|
assert sse.data == ''
|
|
|
|
sse = next(it)
|
|
assert sse.event == 'completion'
|
|
assert sse.data == ''
|
|
|
|
with pytest.raises(StopIteration):
|
|
next(it)
|
|
|
|
|
|
def test_multiple_events_with_data() -> None:
|
|
def body() -> Iterator[str]:
|
|
yield 'event: ping'
|
|
yield 'data: {"foo":true}'
|
|
yield ''
|
|
yield 'event: completion'
|
|
yield 'data: {"bar":false}'
|
|
yield ''
|
|
|
|
it = SSELineParser().iter_lines(body())
|
|
|
|
sse = next(it)
|
|
assert sse.event == 'ping'
|
|
assert sse.json_data() == {'foo': True}
|
|
|
|
sse = next(it)
|
|
assert sse.event == 'completion'
|
|
assert sse.json_data() == {'bar': False}
|
|
|
|
with pytest.raises(StopIteration):
|
|
next(it)
|