82 lines
2.2 KiB
Python
82 lines
2.2 KiB
Python
import json
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from app.ollama import OllamaClient, OllamaError
|
|
|
|
|
|
def make_client(handler, **kwargs):
|
|
return OllamaClient("http://ollama:11434", token="tok", backoff=0,
|
|
transport=httpx.MockTransport(handler), **kwargs)
|
|
|
|
|
|
async def test_chat_payload_and_message():
|
|
seen = {}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
seen["auth"] = request.headers["authorization"]
|
|
seen["body"] = json.loads(request.content)
|
|
return httpx.Response(200, json={"message": {"role": "assistant", "content": "hi"}})
|
|
|
|
msg = await make_client(handler).chat("qwen3", [{"role": "user", "content": "x"}])
|
|
|
|
assert msg["content"] == "hi"
|
|
assert seen["auth"] == "Bearer tok"
|
|
assert seen["body"]["stream"] is False and seen["body"]["model"] == "qwen3"
|
|
assert "tools" not in seen["body"]
|
|
|
|
|
|
async def test_retry_on_5xx_then_success():
|
|
calls = []
|
|
|
|
def handler(request):
|
|
calls.append(1)
|
|
if len(calls) < 3:
|
|
return httpx.Response(503)
|
|
return httpx.Response(200, json={"version": "0.5.0"})
|
|
|
|
assert await make_client(handler).version() == "0.5.0"
|
|
assert len(calls) == 3
|
|
|
|
|
|
async def test_gives_up_after_retries():
|
|
calls = []
|
|
|
|
def handler(request):
|
|
calls.append(1)
|
|
return httpx.Response(500)
|
|
|
|
with pytest.raises(OllamaError, match="500"):
|
|
await make_client(handler).version()
|
|
assert len(calls) == 3
|
|
|
|
|
|
async def test_no_retry_on_4xx():
|
|
calls = []
|
|
|
|
def handler(request):
|
|
calls.append(1)
|
|
return httpx.Response(404, text="model not found")
|
|
|
|
with pytest.raises(OllamaError, match="404"):
|
|
await make_client(handler).chat("nope", [])
|
|
assert len(calls) == 1
|
|
|
|
|
|
async def test_list_models():
|
|
client = make_client(lambda r: httpx.Response(200, json={"models": [{"name": "qwen3"}, {"name": "hermes3"}]}))
|
|
assert await client.list_models() == ["qwen3", "hermes3"]
|
|
|
|
|
|
async def test_transport_error_retried_then_raises():
|
|
calls = []
|
|
|
|
def handler(request):
|
|
calls.append(1)
|
|
raise httpx.ConnectError("down")
|
|
|
|
with pytest.raises(OllamaError, match="ConnectError"):
|
|
await make_client(handler).version()
|
|
assert len(calls) == 3
|