From 21aad8447c55666c9f0a3c7e16f43cd86bae3ca7 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Mon, 22 Jan 2024 14:55:32 -0800 Subject: [PATCH 1/3] fix: update async stream tests --- tests/test_client.py | 57 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 51 insertions(+), 6 deletions(-) diff --git a/tests/test_client.py b/tests/test_client.py index 2f6c8c1..1987531 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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 From c8f7d9b7d271f60ecc16e7ecf9ca098396a11d1e Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Mon, 22 Jan 2024 15:10:18 -0800 Subject: [PATCH 2/3] fix tests --- tests/test_client.py | 54 +++++++++++++++++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 8 deletions(-) diff --git a/tests/test_client.py b/tests/test_client.py index 1987531..c7cb551 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -267,11 +267,14 @@ def test_client_pull_stream(httpserver: HTTPServer): 'insecure': False, 'stream': True, }, - ).respond_with_json({}) + ).respond_with_handler(stream_handler) client = Client(httpserver.url_for('/')) response = client.pull('dummy', stream=True) - assert isinstance(response, types.GeneratorType) + + it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) + for part in response: + assert part['status'] == next(it) def test_client_push(httpserver: HTTPServer): @@ -291,6 +294,14 @@ def test_client_push(httpserver: HTTPServer): def test_client_push_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'retrieving manifest'}) + '\n' + yield json.dumps({'status': 'pushing manifest'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + httpserver.expect_ordered_request( '/api/push', method='POST', @@ -299,11 +310,14 @@ def test_client_push_stream(httpserver: HTTPServer): 'insecure': False, 'stream': True, }, - ).respond_with_json({}) + ).respond_with_handler(stream_handler) client = Client(httpserver.url_for('/')) response = client.push('dummy', stream=True) - assert isinstance(response, types.GeneratorType) + + it = iter(['retrieving manifest', 'pushing manifest', 'success']) + for part in response: + assert part['status'] == next(it) def test_client_create_path(httpserver: HTTPServer): @@ -642,6 +656,16 @@ async def test_async_client_pull(httpserver: HTTPServer): @pytest.mark.asyncio async def test_async_client_pull_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'pulling manifest'}) + '\n' + yield json.dumps({'status': 'verifying sha256 digest'}) + '\n' + yield json.dumps({'status': 'writing manifest'}) + '\n' + yield json.dumps({'status': 'removing any unused layers'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + httpserver.expect_ordered_request( '/api/pull', method='POST', @@ -650,11 +674,14 @@ async def test_async_client_pull_stream(httpserver: HTTPServer): 'insecure': False, 'stream': True, }, - ).respond_with_json({}) + ).respond_with_handler(stream_handler) client = AsyncClient(httpserver.url_for('/')) response = await client.pull('dummy', stream=True) - assert isinstance(response, types.AsyncGeneratorType) + + it = iter(['pulling manifest', 'verifying sha256 digest', 'writing manifest', 'removing any unused layers', 'success']) + async for part in response: + assert part['status'] == next(it) @pytest.mark.asyncio @@ -676,6 +703,14 @@ async def test_async_client_push(httpserver: HTTPServer): @pytest.mark.asyncio async def test_async_client_push_stream(httpserver: HTTPServer): + def stream_handler(_: Request): + def generate(): + yield json.dumps({'status': 'retrieving manifest'}) + '\n' + yield json.dumps({'status': 'pushing manifest'}) + '\n' + yield json.dumps({'status': 'success'}) + '\n' + + return Response(generate()) + httpserver.expect_ordered_request( '/api/push', method='POST', @@ -684,11 +719,14 @@ async def test_async_client_push_stream(httpserver: HTTPServer): 'insecure': False, 'stream': True, }, - ).respond_with_json({}) + ).respond_with_handler(stream_handler) client = AsyncClient(httpserver.url_for('/')) response = await client.push('dummy', stream=True) - assert isinstance(response, types.AsyncGeneratorType) + + it = iter(['retrieving manifest', 'pushing manifest', 'success']) + async for part in response: + assert part['status'] == next(it) @pytest.mark.asyncio From 5c1df78b7dbd63b0e2f037045277f6882bc23084 Mon Sep 17 00:00:00 2001 From: Michael Yang Date: Mon, 22 Jan 2024 15:23:19 -0800 Subject: [PATCH 3/3] remove unused import --- tests/test_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_client.py b/tests/test_client.py index c7cb551..6afbe70 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,7 +1,6 @@ import os import io import json -import types import pytest import tempfile from pathlib import Path