87 lines
2.4 KiB
Python
87 lines
2.4 KiB
Python
|
"""test_embed.py —— ollama 嵌入封装（httpx.Client 注入，不碰真 ollama）。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from vstore import embed
|
||
|
|
from vstore.embed import EmbedError, embed_batch
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeResp:
|
||
|
|
def __init__(self, payload=None, fail=False):
|
||
|
|
self._payload = payload
|
||
|
|
self._fail = fail
|
||
|
|
|
||
|
|
def raise_for_status(self):
|
||
|
|
if self._fail:
|
||
|
|
raise httpx.ConnectError("down", request=None)
|
||
|
|
|
||
|
|
def json(self):
|
||
|
|
return self._payload
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeClient:
|
||
|
|
"""记录调用，返回确定性向量。"""
|
||
|
|
|
||
|
|
def __init__(self, vectors=None, fail=False):
|
||
|
|
self.vectors = vectors or []
|
||
|
|
self.fail = fail
|
||
|
|
self.calls = []
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *args):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def post(self, url, json=None, timeout=None):
|
||
|
|
self.calls.append((url, json))
|
||
|
|
if self.fail:
|
||
|
|
return _FakeResp(fail=True)
|
||
|
|
text = (json or {}).get("input", "")
|
||
|
|
return _FakeResp({"embeddings": [[1.0] + [0.0] * 1023]})
|
||
|
|
|
||
|
|
|
||
|
|
def _vec():
|
||
|
|
return [1.0] + [0.0] * 1023
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed_batch_ok(monkeypatch):
|
||
|
|
fake = _FakeClient()
|
||
|
|
monkeypatch.setattr(embed.httpx, "Client", lambda: fake)
|
||
|
|
out = embed_batch(["hello", "world"], url="http://x", model="m")
|
||
|
|
assert len(out) == 2
|
||
|
|
assert out[0] == _vec()
|
||
|
|
assert len(fake.calls) == 2
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed_batch_empty_text_no_upstream(monkeypatch):
|
||
|
|
fake = _FakeClient()
|
||
|
|
monkeypatch.setattr(embed.httpx, "Client", lambda: fake)
|
||
|
|
out = embed_batch(["", " ", "hello"])
|
||
|
|
assert len(out) == 3
|
||
|
|
assert out[0] == [0.0] * 1024 # 空文本零向量
|
||
|
|
assert len(fake.calls) == 1 # 只有 "hello" 打上游
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed_batch_upstream_fail_degrades_zero(monkeypatch):
|
||
|
|
fake = _FakeClient(fail=True)
|
||
|
|
monkeypatch.setattr(embed.httpx, "Client", lambda: fake)
|
||
|
|
out = embed_batch(["hello"])
|
||
|
|
assert out == [[0.0] * 1024] # 降级零向量，不抛
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed_batch_bad_dim_degrades(monkeypatch):
|
||
|
|
class BadClient(_FakeClient):
|
||
|
|
def post(self, url, json=None, timeout=None):
|
||
|
|
return _FakeResp({"embeddings": [[1.0, 2.0]]}) # 维度不符
|
||
|
|
|
||
|
|
monkeypatch.setattr(embed.httpx, "Client", lambda: BadClient())
|
||
|
|
out = embed_batch(["hello"])
|
||
|
|
assert out == [[0.0] * 1024]
|
||
|
|
|
||
|
|
|
||
|
|
def test_embed_batch_empty_list():
|
||
|
|
assert embed_batch([]) == []
|