Spaces:
Running
Running
File size: 3,907 Bytes
1f5b714 c449c99 1f5b714 c449c99 1f5b714 c449c99 1f5b714 c449c99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
"""Unit tests for RAG Index Service."""
import pytest
from unittest.mock import MagicMock, patch
from pathlib import Path
from backend.src.services.rag_index import RAGIndexService
from backend.src.services.config import AppConfig
@pytest.fixture
def mock_config(tmp_path):
config = MagicMock(spec=AppConfig)
config.google_api_key = "fake-key"
config.llamaindex_persist_dir = tmp_path / "data" / "llamaindex"
return config
@pytest.fixture
def rag_service(mock_config):
with patch("backend.src.services.rag_index.get_config", return_value=mock_config):
with patch("backend.src.services.rag_index.VaultService") as mock_vault:
service = RAGIndexService()
service.vault_service = mock_vault.return_value
yield service
def test_get_persist_dir(rag_service):
user_id = "test-user"
persist_dir = rag_service.get_persist_dir(user_id)
assert user_id in persist_dir
assert Path(persist_dir).exists()
@patch("backend.src.services.rag_index.load_index_from_storage")
@patch("backend.src.services.rag_index.StorageContext")
def test_get_or_build_index_existing(mock_storage_context, mock_load, rag_service):
"""Test loading an existing index."""
user_id = "test-user"
mock_index = MagicMock()
mock_load.return_value = mock_index
# Assume load succeeds
index = rag_service.get_or_build_index(user_id)
assert index == mock_index
mock_load.assert_called_once()
# Should NOT build
rag_service.vault_service.list_notes.assert_not_called()
@patch("backend.src.services.rag_index.VectorStoreIndex")
@patch("backend.src.services.rag_index.load_index_from_storage")
def test_get_or_build_index_new(mock_load, mock_vector_store, rag_service):
"""Test building a new index when load fails."""
user_id = "test-user"
mock_load.side_effect = Exception("No index")
mock_index = MagicMock()
mock_vector_store.from_documents.return_value = mock_index
# Mock vault data
rag_service.vault_service.list_notes.return_value = [{"path": "note1.md"}]
rag_service.vault_service.read_note.return_value = {
"title": "Note 1",
"body": "Content",
"metadata": {}
}
index = rag_service.get_or_build_index(user_id)
rag_service.vault_service.list_notes.assert_called_with(user_id)
@patch("backend.src.services.rag_index.os.path.exists")
def test_get_status(mock_exists, rag_service):
user_id = "test-user"
mock_exists.return_value = True
status = rag_service.get_status(user_id)
assert status.status == "ready"
mock_exists.return_value = False
status = rag_service.get_status(user_id)
assert status.status == "building"
@patch("backend.src.services.rag_index.load_index_from_storage")
@patch("backend.src.services.rag_index.StorageContext")
def test_chat(mock_storage, mock_load, rag_service):
user_id = "test-user"
# Mock Index and ChatEngine
mock_index = MagicMock()
mock_chat_engine = MagicMock()
mock_index.as_chat_engine.return_value = mock_chat_engine
mock_load.return_value = mock_index
# Mock Response
mock_response = MagicMock()
mock_response.__str__.return_value = "AI Answer"
# Mock Source Nodes
mock_node = MagicMock()
mock_node.metadata = {"path": "note.md", "title": "Note"}
mock_node.get_content.return_value = "Snippet content"
mock_node.score = 0.9
mock_response.source_nodes = [mock_node]
mock_chat_engine.chat.return_value = mock_response
from backend.src.models.rag import ChatMessage
messages = [ChatMessage(role="user", content="Question")]
response = rag_service.chat(user_id, messages)
assert response.answer == "AI Answer"
assert len(response.sources) == 1
assert response.sources[0].path == "note.md"
mock_chat_engine.chat.assert_called()
|