forked from HSE_team/BetterCallPraskovia
171 lines
5.3 KiB
Python
171 lines
5.3 KiB
Python
import pytest
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
from uuid import uuid4
|
|
from datetime import datetime
|
|
from typing import List, Dict
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_document():
|
|
try:
|
|
from backend.src.domain.entities.document import Document
|
|
except ImportError:
|
|
from src.domain.entities.document import Document
|
|
return Document(
|
|
collection_id=uuid4(),
|
|
title="Тестовый документ",
|
|
content="Содержание документа",
|
|
metadata={"type": "law", "article": "123"}
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_documents_list():
|
|
try:
|
|
from backend.src.domain.entities.document import Document
|
|
except ImportError:
|
|
from src.domain.entities.document import Document
|
|
collection_id = uuid4()
|
|
return [
|
|
Document(
|
|
collection_id=collection_id,
|
|
title=f"Документ {i}",
|
|
content=f"Содержание документа {i} ",
|
|
metadata={"relevance_score": 0.9 - i * 0.1}
|
|
)
|
|
for i in range(10)
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_relevant_documents():
|
|
return [
|
|
{"document_id": str(uuid4()), "title": "Гражданский кодекс РФ", "relevance": True},
|
|
{"document_id": str(uuid4()), "title": "Трудовой кодекс РФ", "relevance": True},
|
|
{"document_id": str(uuid4()), "title": "Налоговый кодекс РФ", "relevance": True},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_rag_response():
|
|
return {
|
|
"answer": "Тестовый ответ на вопрос",
|
|
"sources": [
|
|
{"title": "Документ 1", "collection": "Коллекция 1", "document_id": str(uuid4())},
|
|
{"title": "Документ 2", "collection": "Коллекция 1", "document_id": str(uuid4())},
|
|
{"title": "Документ 3", "collection": "Коллекция 2", "document_id": str(uuid4())},
|
|
],
|
|
"usage": {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300}
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_collection():
|
|
try:
|
|
from backend.src.domain.entities.collection import Collection
|
|
except ImportError:
|
|
from src.domain.entities.collection import Collection
|
|
return Collection(
|
|
name="Коллекция",
|
|
owner_id=uuid4(),
|
|
description="Описание коллекции",
|
|
is_public=False
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_user():
|
|
try:
|
|
from backend.src.domain.entities.user import User, UserRole
|
|
except ImportError:
|
|
from src.domain.entities.user import User, UserRole
|
|
return User(
|
|
telegram_id="123456789",
|
|
role=UserRole.USER
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_document_repository():
|
|
repository = AsyncMock()
|
|
repository.get_by_id = AsyncMock()
|
|
repository.create = AsyncMock()
|
|
repository.update = AsyncMock()
|
|
repository.delete = AsyncMock(return_value=True)
|
|
repository.list_by_collection = AsyncMock(return_value=[])
|
|
return repository
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_collection_repository():
|
|
repository = AsyncMock()
|
|
repository.get_by_id = AsyncMock()
|
|
repository.create = AsyncMock()
|
|
repository.update = AsyncMock()
|
|
repository.delete = AsyncMock(return_value=True)
|
|
repository.list_by_owner = AsyncMock(return_value=[])
|
|
repository.list_public = AsyncMock(return_value=[])
|
|
return repository
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_user_repository():
|
|
repository = AsyncMock()
|
|
repository.get_by_id = AsyncMock()
|
|
repository.get_by_telegram_id = AsyncMock()
|
|
repository.create = AsyncMock()
|
|
repository.update = AsyncMock()
|
|
repository.delete = AsyncMock(return_value=True)
|
|
return repository
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_deepseek_client():
|
|
client = AsyncMock()
|
|
client.chat_completion = AsyncMock(return_value={
|
|
"content": "Ответ от DeepSeek",
|
|
"usage": {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300}
|
|
})
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_queries_with_ground_truth():
|
|
return [
|
|
{
|
|
"query": "Какие права имеет работник при увольнении?",
|
|
"relevant_document_ids": [str(uuid4()), str(uuid4())],
|
|
"expected_top5_contains": True
|
|
},
|
|
{
|
|
"query": "Как оформить договор купли-продажи?",
|
|
"relevant_document_ids": [str(uuid4())],
|
|
"expected_top5_contains": True
|
|
},
|
|
{
|
|
"query": "Какие налоги платит ИП?",
|
|
"relevant_document_ids": [str(uuid4()), str(uuid4()), str(uuid4())],
|
|
"expected_top5_contains": True
|
|
},
|
|
{
|
|
"query": "Права потребителя при возврате товара",
|
|
"relevant_document_ids": [str(uuid4())],
|
|
"expected_top5_contains": True
|
|
},
|
|
{
|
|
"query": "Как расторгнуть брак?",
|
|
"relevant_document_ids": [str(uuid4())],
|
|
"expected_top5_contains": True
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_aiohttp_session():
|
|
session = AsyncMock()
|
|
session.get = AsyncMock()
|
|
session.post = AsyncMock()
|
|
session.__aenter__ = AsyncMock(return_value=session)
|
|
session.__aexit__ = AsyncMock(return_value=None)
|
|
return session
|