feat: ollama embedder mit exponential backoff retry
This commit is contained in:
52
tests/test_embedder.py
Normal file
52
tests/test_embedder.py
Normal file
@@ -0,0 +1,52 @@
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
import pytest
|
||||
|
||||
from app.ingest.embedder import embed_texts, EmbeddingError
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embed_texts_returns_vectors():
|
||||
fake_client = MagicMock()
|
||||
fake_client.embeddings = AsyncMock(side_effect=[
|
||||
{"embedding": [0.1, 0.2, 0.3]},
|
||||
{"embedding": [0.4, 0.5, 0.6]},
|
||||
])
|
||||
|
||||
with patch("app.ingest.embedder._client", return_value=fake_client):
|
||||
vectors = await embed_texts(["hello", "world"], model="qwen3-embedding:0.6b")
|
||||
|
||||
assert vectors == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
assert fake_client.embeddings.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embed_texts_retries_on_failure(monkeypatch):
|
||||
monkeypatch.setattr("app.ingest.embedder._BACKOFF_SECONDS", [0, 0, 0])
|
||||
|
||||
fake_client = MagicMock()
|
||||
fake_client.embeddings = AsyncMock(side_effect=[
|
||||
Exception("connection refused"),
|
||||
Exception("timeout"),
|
||||
{"embedding": [1.0, 2.0]},
|
||||
])
|
||||
|
||||
with patch("app.ingest.embedder._client", return_value=fake_client):
|
||||
vectors = await embed_texts(["hi"], model="m")
|
||||
|
||||
assert vectors == [[1.0, 2.0]]
|
||||
assert fake_client.embeddings.call_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_embed_texts_raises_after_max_retries(monkeypatch):
|
||||
monkeypatch.setattr("app.ingest.embedder._BACKOFF_SECONDS", [0, 0, 0])
|
||||
|
||||
fake_client = MagicMock()
|
||||
fake_client.embeddings = AsyncMock(side_effect=Exception("nope"))
|
||||
|
||||
with patch("app.ingest.embedder._client", return_value=fake_client):
|
||||
with pytest.raises(EmbeddingError):
|
||||
await embed_texts(["hi"], model="m")
|
||||
|
||||
assert fake_client.embeddings.call_count == 4 # initial + 3 retries
|
||||
Reference in New Issue
Block a user