fix: update async stream tests

This commit is contained in:
Michael Yang 2024-01-22 14:55:32 -08:00
parent 38f68e251d
commit 21aad8447c

View File

@ -81,9 +81,11 @@ def test_client_chat_stream(httpserver: HTTPServer):
client = Client(httpserver.url_for('/'))
response = client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True)
it = iter(['I ', "don't ", 'know.'])
for part in response:
assert part['message']['role'] in 'assistant'
assert part['message']['content'] in ['I ', "don't ", 'know.']
assert part['message']['content'] == next(it)
def test_client_chat_images(httpserver: HTTPServer):
@ -187,9 +189,11 @@ def test_client_generate_stream(httpserver: HTTPServer):
client = Client(httpserver.url_for('/'))
response = client.generate('dummy', 'Why is the sky blue?', stream=True)
it = iter(['Because ', 'it ', 'is.'])
for part in response:
assert part['model'] == 'dummy'
assert part['response'] in ['Because ', 'it ', 'is.']
assert part['response'] == next(it)
def test_client_generate_images(httpserver: HTTPServer):
@ -458,6 +462,24 @@ async def test_async_client_chat(httpserver: HTTPServer):
@pytest.mark.asyncio
async def test_async_client_chat_stream(httpserver: HTTPServer):
def stream_handler(_: Request):
def generate():
for message in ['I ', "don't ", 'know.']:
yield (
json.dumps(
{
'model': 'dummy',
'message': {
'role': 'assistant',
'content': message,
},
}
)
+ '\n'
)
return Response(generate())
httpserver.expect_ordered_request(
'/api/chat',
method='POST',
@ -468,11 +490,15 @@ async def test_async_client_chat_stream(httpserver: HTTPServer):
'format': '',
'options': {},
},
).respond_with_json({})
).respond_with_handler(stream_handler)
client = AsyncClient(httpserver.url_for('/'))
response = await client.chat('dummy', messages=[{'role': 'user', 'content': 'Why is the sky blue?'}], stream=True)
assert isinstance(response, types.AsyncGeneratorType)
it = iter(['I ', "don't ", 'know.'])
async for part in response:
assert part['message']['role'] == 'assistant'
assert part['message']['content'] == next(it)
@pytest.mark.asyncio
@ -529,6 +555,21 @@ async def test_async_client_generate(httpserver: HTTPServer):
@pytest.mark.asyncio
async def test_async_client_generate_stream(httpserver: HTTPServer):
def stream_handler(_: Request):
def generate():
for message in ['Because ', 'it ', 'is.']:
yield (
json.dumps(
{
'model': 'dummy',
'response': message,
}
)
+ '\n'
)
return Response(generate())
httpserver.expect_ordered_request(
'/api/generate',
method='POST',
@ -544,11 +585,15 @@ async def test_async_client_generate_stream(httpserver: HTTPServer):
'format': '',
'options': {},
},
).respond_with_json({})
).respond_with_handler(stream_handler)
client = AsyncClient(httpserver.url_for('/'))
response = await client.generate('dummy', 'Why is the sky blue?', stream=True)
assert isinstance(response, types.AsyncGeneratorType)
it = iter(['Because ', 'it ', 'is.'])
async for part in response:
assert part['model'] == 'dummy'
assert part['response'] == next(it)
@pytest.mark.asyncio