kanshan/ollama-python
1
0
Fork 0
mirror of https://github.com/ollama/ollama-python.git synced 2026-01-13 21:57:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix tests

Browse Source
This commit is contained in:
Michael Yang 2024-01-22 15:10:18 -08:00
parent 21aad8447c
commit c8f7d9b7d2
1 changed files with 46 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

54
tests/test_client.py
View File

@ -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