656f596d7e
Core modules: - Laws: CRUD, search, AI-powered QA - Analysis: legal research and case management - Contracts: lifecycle management with templates - Signatures: electronic signature workflow Infrastructure: - FastAPI + SQLite + async SQLAlchemy - Docker deployment support - 54 unit tests passing Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
"""Unit tests for Vector service."""
|
|
import pytest
|
|
import numpy as np
|
|
from unittest.mock import AsyncMock, patch, MagicMock
|
|
|
|
from app.services.vector_service import VectorService
|
|
|
|
|
|
class TestVectorService:
|
|
"""Test cases for Vector service."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_embedding_no_api_key(self):
|
|
"""Test getting embedding without API key returns mock."""
|
|
service = VectorService()
|
|
service.api_key = None # Ensure no API key
|
|
|
|
embedding = await service.get_embedding("测试文本")
|
|
|
|
# Should return mock embedding
|
|
assert len(embedding) == service.dimension
|
|
assert all(x == 0.0 for x in embedding)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_embedding_with_api_key(self):
|
|
"""Test getting embedding with API key."""
|
|
service = VectorService()
|
|
service.api_key = "test-key"
|
|
|
|
with patch('httpx.AsyncClient') as mock_client:
|
|
mock_response = MagicMock()
|
|
mock_response.json.return_value = {
|
|
"data": [{"embedding": [0.1] * 1536}]
|
|
}
|
|
mock_response.raise_for_status = MagicMock()
|
|
|
|
mock_client_instance = MagicMock()
|
|
mock_client_instance.post = AsyncMock(return_value=mock_response)
|
|
mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance)
|
|
mock_client_instance.__aexit__ = AsyncMock(return_value=None)
|
|
mock_client.return_value = mock_client_instance
|
|
|
|
embedding = await service.get_embedding("测试文本")
|
|
|
|
assert len(embedding) == 1536
|
|
assert all(x == 0.1 for x in embedding)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cosine_similarity(self):
|
|
"""Test cosine similarity calculation."""
|
|
service = VectorService()
|
|
|
|
vec1 = [1.0, 0.0, 0.0]
|
|
vec2 = [1.0, 0.0, 0.0]
|
|
vec3 = [0.0, 1.0, 0.0]
|
|
|
|
# Same vectors should have similarity 1.0
|
|
sim1 = service.cosine_similarity(vec1, vec2)
|
|
assert abs(sim1 - 1.0) < 0.001
|
|
|
|
# Orthogonal vectors should have similarity 0.0
|
|
sim2 = service.cosine_similarity(vec1, vec3)
|
|
assert abs(sim2 - 0.0) < 0.001
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_search_similar(self):
|
|
"""Test searching similar vectors."""
|
|
service = VectorService()
|
|
|
|
# Mock vectors: first is similar, second is different
|
|
stored_vectors = [
|
|
{"id": 1, "embedding": [1.0, 0.0, 0.0]},
|
|
{"id": 2, "embedding": [0.0, 1.0, 0.0]},
|
|
]
|
|
|
|
query_vec = [0.9, 0.1, 0.0]
|
|
|
|
results = service.search_similar(query_vec, stored_vectors, top_k=2)
|
|
|
|
assert len(results) == 2
|
|
assert results[0]["id"] == 1 # Most similar should be first
|