import pytest
from unittest.mock import MagicMock, patch
from sqlalchemy.orm import Session
from sqlalchemy.exc import SQLAlchemyError
from datetime import datetime
from app.core.services.document import DocumentService
from app.db import models
from app.core.vector_store.faiss_store import FaissVectorStore
from app.core.vector_store.embedder.mock import MockEmbedder
@pytest.fixture
def document_service():
"""
Pytest fixture to create a DocumentService instance with mocked dependencies.
"""
mock_embedder = MagicMock(spec=MockEmbedder)
mock_vector_store = MagicMock(spec=FaissVectorStore)
mock_vector_store.embedder = mock_embedder
return DocumentService(vector_store=mock_vector_store)
# --- add_document Tests ---
def test_add_document_success(document_service: DocumentService):
"""
Tests that add_document successfully adds a document to the database
and its vector embedding to the FAISS index.
"""
# Arrange
mock_db = MagicMock(spec=Session)
mock_document = MagicMock(id=1, text="Test text.")
mock_db.add.side_effect = [None, None] # Allow multiple calls
# Configure the mock db.query to return a document object
mock_document_model_instance = models.Document(
id=1,
title="Test Title",
text="Test text.",
source_url="http://test.com"
)
with patch('app.core.services.document.models.Document', return_value=mock_document_model_instance) as mock_document_model, \
patch('app.core.services.document.models.VectorMetadata') as mock_vector_metadata_model:
document_service.vector_store.add_document.return_value = 123
doc_data = {
"title": "Test Title",
"text": "Test text.",
"source_url": "http://test.com"
}
# Act
document_id = document_service.add_document(db=mock_db, doc_data=doc_data)
# Assert
assert document_id == 1
mock_db.add.assert_any_call(mock_document_model_instance)
mock_db.commit.assert_called()
mock_db.refresh.assert_called_with(mock_document_model_instance)
document_service.vector_store.add_document.assert_called_once_with(text="Test text.",document_id=1,db_session=mock_db,embedding_model="mock_embedder")
# mock_vector_metadata_model.assert_called_once_with(
# document_id=1,
# faiss_index=123,
# embedding_model="mock_embedder"
# )
# mock_db.add.assert_any_call(mock_vector_metadata_model.return_value)
def test_add_document_sql_error(document_service: DocumentService):
"""
Tests that add_document correctly handles a SQLAlchemyError by rolling back.
"""
# Arrange
mock_db = MagicMock(spec=Session)
mock_db.add.side_effect = SQLAlchemyError("Database error")
doc_data = {"title": "Test", "text": "...", "source_url": "http://test.com"}
# Act & Assert
with pytest.raises(SQLAlchemyError, match="Database error"):
document_service.add_document(db=mock_db, doc_data=doc_data)
mock_db.rollback.assert_called_once()
mock_db.commit.assert_not_called()
# --- get_all_documents Tests ---
def test_get_all_documents_success(document_service: DocumentService):
"""
Tests that get_all_documents returns a list of documents.
"""
# Arrange
mock_db = MagicMock(spec=Session)
mock_documents = [models.Document(id=1), models.Document(id=2)]
mock_db.query.return_value.order_by.return_value.all.return_value = mock_documents
# Act
documents = document_service.get_all_documents(db=mock_db)
# Assert
assert documents == mock_documents
mock_db.query.assert_called_once_with(models.Document)
mock_db.query.return_value.order_by.assert_called_once()
# --- delete_document Tests ---
def test_delete_document_success(document_service: DocumentService):
"""
Tests that delete_document correctly deletes a document.
"""
# Arrange
mock_db = MagicMock(spec=Session)
doc_id_to_delete = 1
doc_to_delete = models.Document(id=doc_id_to_delete)
mock_db.query.return_value.filter.return_value.first.return_value = doc_to_delete
# Act
deleted_id = document_service.delete_document(db=mock_db, document_id=doc_id_to_delete)
# Assert
assert deleted_id == doc_id_to_delete
mock_db.query.assert_called_once_with(models.Document)
mock_db.delete.assert_called_once_with(doc_to_delete)
mock_db.commit.assert_called_once()
def test_delete_document_not_found(document_service: DocumentService):
"""
Tests that delete_document returns None if the document is not found.
"""
# Arrange
mock_db = MagicMock(spec=Session)
mock_db.query.return_value.filter.return_value.first.return_value = None
# Act
deleted_id = document_service.delete_document(db=mock_db, document_id=999)
# Assert
assert deleted_id is None
mock_db.delete.assert_not_called()
mock_db.commit.assert_not_called()
def test_delete_document_sql_error(document_service: DocumentService):
"""
Tests that delete_document handles a SQLAlchemyError correctly by rolling back.
"""
# Arrange
mock_db = MagicMock(spec=Session)
doc_id = 1
doc_to_delete = models.Document(id=doc_id)
mock_db.query.return_value.filter.return_value.first.return_value = doc_to_delete
mock_db.delete.side_effect = SQLAlchemyError("Delete error")
# Act & Assert
with pytest.raises(SQLAlchemyError, match="Delete error"):
document_service.delete_document(db=mock_db, document_id=doc_id)
mock_db.rollback.assert_called_once()
mock_db.commit.assert_not_called()