Newer
Older
cortex-hub / ai-hub / tests / core / services / test_document.py
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("Test text.")
        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()