kplam
51dc8dc4f6
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索（dense + sparse RRF 融合） - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
"""OllamaClient 单元测试（mock httpx，不发起真实网络请求）"""
|
||
|
||
from typing import Any
|
||
|
||
import httpx
|
||
import pytest
|
||
|
||
from app.services.ollama import OllamaClient
|
||
|
||
|
||
def _make_response(status_code: int, data: dict | None = None) -> httpx.Response:
|
||
"""构造带 request 上下文的 httpx.Response（raise_for_status 依赖 request）"""
|
||
request = httpx.Request("POST", "http://localhost:11434/api/generate")
|
||
if data is None:
|
||
return httpx.Response(status_code, request=request)
|
||
return httpx.Response(status_code, json=data, request=request)
|
||
|
||
|
||
class _FakeAsyncClient:
|
||
"""httpx.AsyncClient 替代品：记录请求参数，返回预设响应或抛出预设异常"""
|
||
|
||
def __init__(
|
||
self,
|
||
captured: dict,
|
||
response: httpx.Response | None = None,
|
||
error: Exception | None = None,
|
||
**kwargs: Any,
|
||
) -> None:
|
||
self._captured = captured
|
||
self._response = response
|
||
self._error = error
|
||
captured["timeout"] = kwargs.get("timeout")
|
||
|
||
async def __aenter__(self) -> "_FakeAsyncClient":
|
||
return self
|
||
|
||
async def __aexit__(self, *args: object) -> bool:
|
||
return False
|
||
|
||
async def post(self, url: str, json: dict | None = None) -> httpx.Response:
|
||
self._captured["url"] = url
|
||
self._captured["json"] = json
|
||
if self._error is not None:
|
||
raise self._error
|
||
assert self._response is not None
|
||
return self._response
|
||
|
||
async def get(self, url: str) -> httpx.Response:
|
||
self._captured["url"] = url
|
||
if self._error is not None:
|
||
raise self._error
|
||
assert self._response is not None
|
||
return self._response
|
||
|
||
|
||
class TestGenerate:
|
||
async def test_returns_response_field(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||
captured: dict = {}
|
||
monkeypatch.setattr(
|
||
httpx,
|
||
"AsyncClient",
|
||
lambda **kw: _FakeAsyncClient(captured, response=_make_response(200, {"response": "生成结果"}), **kw),
|
||
)
|
||
client = OllamaClient(base_url="http://localhost:11434/", model="qwen2.5:1.5b")
|
||
|
||
result = await client.generate("你好")
|
||
|
||
assert result == "生成结果"
|
||
# base_url 尾部斜杠被去除
|
||
assert captured["url"] == "http://localhost:11434/api/generate"
|
||
assert captured["json"] == {"model": "qwen2.5:1.5b", "prompt": "你好", "stream": False}
|
||
assert "format" not in captured["json"]
|
||
assert captured["timeout"] == 120.0
|
||
|
||
async def test_json_mode_adds_format(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||
captured: dict = {}
|
||
monkeypatch.setattr(
|
||
httpx,
|
||
"AsyncClient",
|
||
lambda **kw: _FakeAsyncClient(captured, response=_make_response(200, {"response": "{}"}), **kw),
|
||
)
|
||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||
|
||
await client.generate("你好", json_mode=True)
|
||
|
||
assert captured["json"]["format"] == "json"
|
||
|
||
async def test_http_error_status_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(
|
||
httpx, "AsyncClient", lambda **kw: _FakeAsyncClient({}, response=_make_response(500), **kw)
|
||
)
|
||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||
|
||
with pytest.raises(httpx.HTTPStatusError):
|
||
await client.generate("你好")
|
||
|
||
|
||
class TestIsAvailable:
|
||
async def test_200_returns_true(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||
captured: dict = {}
|
||
monkeypatch.setattr(
|
||
httpx,
|
||
"AsyncClient",
|
||
lambda **kw: _FakeAsyncClient(captured, response=_make_response(200, {"models": []}), **kw),
|
||
)
|
||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||
|
||
assert await client.is_available() is True
|
||
assert captured["url"] == "http://localhost:11434/api/tags"
|
||
assert captured["timeout"] == 5.0
|
||
|
||
async def test_non_200_returns_false(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(
|
||
httpx, "AsyncClient", lambda **kw: _FakeAsyncClient({}, response=_make_response(503), **kw)
|
||
)
|
||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||
|
||
assert await client.is_available() is False
|
||
|
||
async def test_connect_error_returns_false(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||
monkeypatch.setattr(
|
||
httpx, "AsyncClient", lambda **kw: _FakeAsyncClient({}, error=httpx.ConnectError("连接被拒绝"), **kw)
|
||
)
|
||
client = OllamaClient(base_url="http://localhost:11434", model="qwen2.5:1.5b")
|
||
|
||
assert await client.is_available() is False
|