""" Tests for document deletion functionality. Tests the DocumentDeletionService methods: - delete_document - delete_blob_only - remove_from_collection """ import pytest import uuid from sqlalchemy import create_engine, event from sqlalchemy.orm import sessionmaker from local_deep_research.database.models.library import ( Document, DocumentBlob, DocumentChunk, DocumentCollection, Collection, SourceType, ) from local_deep_research.constants import ( FILE_PATH_BLOB_DELETED, ) from local_deep_research.database.models.base import Base @pytest.fixture def test_engine(): """Create an in-memory SQLite database for testing.""" engine = create_engine("sqlite:///:memory:") # Enable foreign key enforcement for SQLite (required for CASCADE to work) @event.listens_for(engine, "connect") def set_sqlite_pragma(dbapi_connection, connection_record): cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") cursor.close() Base.metadata.create_all(engine) yield engine engine.dispose() @pytest.fixture def test_session(test_engine): """Create a session for testing.""" Session = sessionmaker(bind=test_engine) session = Session() yield session session.close() @pytest.fixture def source_type(test_session): """Create a source type for testing.""" st = SourceType( id=str(uuid.uuid4()), name="user_upload", display_name="User Upload", ) test_session.add(st) test_session.commit() return st @pytest.fixture def test_collection(test_session): """Create a test collection.""" collection = Collection( id=str(uuid.uuid4()), name="Test Collection", description="A test collection", ) test_session.add(collection) test_session.commit() return collection @pytest.fixture def test_document(test_session, source_type, test_collection): """Create a test document with blob and collection link.""" doc_id = str(uuid.uuid4()) doc = Document( id=doc_id, source_type_id=source_type.id, document_hash=f"hash_{doc_id[:8]}", file_size=1024, file_type="pdf", title="Test Document", filename="test.pdf", storage_mode="database", text_content="This is test content.", ) test_session.add(doc) test_session.flush() # Add blob blob = DocumentBlob( document_id=doc_id, pdf_binary=b"fake pdf content", blob_hash="fake_hash", ) test_session.add(blob) # Add to collection doc_collection = DocumentCollection( document_id=doc_id, collection_id=test_collection.id, indexed=False, ) test_session.add(doc_collection) test_session.commit() return doc class TestDocumentDeletionService: """Unit tests for DocumentDeletionService.""" def test_delete_document_removes_blob(self, test_session, test_document): """Verify DocumentBlob is deleted with document.""" doc_id = test_document.id # Verify blob exists before deletion blob_before = ( test_session.query(DocumentBlob) .filter_by(document_id=doc_id) .first() ) assert blob_before is not None, "Blob should exist before deletion" # Delete document (manually simulating service behavior) doc = test_session.query(Document).get(doc_id) test_session.delete(doc) test_session.commit() # Verify blob is deleted (CASCADE should handle this) blob_after = ( test_session.query(DocumentBlob) .filter_by(document_id=doc_id) .first() ) assert blob_after is None, "Blob should be deleted with document" def test_delete_document_removes_collection_links( self, test_session, test_document, test_collection ): """Verify DocumentCollection links are deleted.""" doc_id = test_document.id # Verify link exists before deletion link_before = ( test_session.query(DocumentCollection) .filter_by(document_id=doc_id) .first() ) assert link_before is not None, "Collection link should exist" # Delete document doc = test_session.query(Document).get(doc_id) test_session.delete(doc) test_session.commit() # Verify link is deleted link_after = ( test_session.query(DocumentCollection) .filter_by(document_id=doc_id) .first() ) assert link_after is None, "Collection link should be deleted" def test_delete_blob_keeps_document(self, test_session, test_document): """Verify document metadata and text_content preserved after blob deletion.""" doc_id = test_document.id # Delete only the blob blob = ( test_session.query(DocumentBlob) .filter_by(document_id=doc_id) .first() ) test_session.delete(blob) test_session.commit() # Verify document still exists with text doc = test_session.query(Document).get(doc_id) assert doc is not None, "Document should still exist" assert doc.text_content is not None, "Text content should be preserved" assert doc.title == "Test Document", "Title should be preserved" def test_remove_from_collection_deletes_orphan( self, test_session, test_document, test_collection ): """Document in single collection: deleted completely.""" doc_id = test_document.id # Remove the collection link (simulate remove_from_collection) link = ( test_session.query(DocumentCollection) .filter_by(document_id=doc_id, collection_id=test_collection.id) .first() ) test_session.delete(link) test_session.flush() # Check if document is in any collection remaining = ( test_session.query(DocumentCollection) .filter_by(document_id=doc_id) .count() ) assert remaining == 0, "Document should not be in any collection" # If orphaned, delete document if remaining == 0: doc = test_session.query(Document).get(doc_id) test_session.delete(doc) test_session.commit() # Verify document is deleted doc_after = test_session.query(Document).get(doc_id) assert doc_after is None, "Orphaned document should be deleted" def test_remove_from_collection_keeps_if_in_other( self, test_session, test_document, test_collection ): """Document in multiple collections: only unlinked.""" doc_id = test_document.id # Create another collection and link document to it other_collection = Collection( id=str(uuid.uuid4()), name="Other Collection", ) test_session.add(other_collection) test_session.flush() other_link = DocumentCollection( document_id=doc_id, collection_id=other_collection.id, indexed=False, ) test_session.add(other_link) test_session.commit() # Remove from first collection link = ( test_session.query(DocumentCollection) .filter_by(document_id=doc_id, collection_id=test_collection.id) .first() ) test_session.delete(link) test_session.commit() # Document should still exist (in other collection) doc = test_session.query(Document).get(doc_id) assert doc is not None, "Document should exist (in other collection)" # Document should only be in other collection links = ( test_session.query(DocumentCollection) .filter_by(document_id=doc_id) .all() ) assert len(links) == 1, "Document should be in exactly one collection" assert links[0].collection_id == other_collection.id class TestBlobDeletion: """Unit tests for blob-only deletion.""" def test_delete_blob_updates_storage_mode( self, test_session, test_document ): """Verify storage_mode set to 'none' after blob deletion.""" doc_id = test_document.id # Delete blob and update storage mode blob = ( test_session.query(DocumentBlob) .filter_by(document_id=doc_id) .first() ) test_session.delete(blob) doc = test_session.query(Document).get(doc_id) doc.storage_mode = "none" doc.file_path = FILE_PATH_BLOB_DELETED test_session.commit() # Verify update doc_after = test_session.query(Document).get(doc_id) assert doc_after.storage_mode == "none" assert doc_after.file_path == FILE_PATH_BLOB_DELETED def test_delete_blob_frees_bytes(self, test_session, test_document): """Verify bytes freed calculation is correct.""" doc_id = test_document.id # Get blob size blob = ( test_session.query(DocumentBlob) .filter_by(document_id=doc_id) .first() ) blob_size = len(blob.pdf_binary) if blob.pdf_binary else 0 assert blob_size > 0, "Blob should have content" # Delete blob test_session.delete(blob) test_session.commit() # Verify blob is deleted blob_after = ( test_session.query(DocumentBlob) .filter_by(document_id=doc_id) .first() ) assert blob_after is None class TestDocumentChunkCleanup: """Tests for DocumentChunk cleanup (no FK constraint).""" def test_chunks_require_manual_cleanup( self, test_session, test_document, test_collection ): """Verify DocumentChunks must be manually deleted (no FK).""" doc_id = test_document.id collection_name = f"collection_{test_collection.id}" # Create some chunks for the document for i in range(3): chunk = DocumentChunk( chunk_hash=f"chunk_hash_{i}_{uuid.uuid4().hex[:8]}", source_type="document", source_id=doc_id, collection_name=collection_name, chunk_text=f"Test chunk {i} content", chunk_index=i, start_char=i * 100, end_char=(i + 1) * 100, word_count=10, embedding_id=str(uuid.uuid4()), embedding_model="test-model", embedding_model_type="sentence_transformers", ) test_session.add(chunk) test_session.commit() # Verify chunks exist chunks_before = ( test_session.query(DocumentChunk) .filter_by(source_id=doc_id, source_type="document") .count() ) assert chunks_before == 3, "Should have 3 chunks" # Delete document WITHOUT manually deleting chunks first doc = test_session.query(Document).get(doc_id) test_session.delete(doc) test_session.commit() # Chunks should still exist (no FK constraint!) chunks_after = ( test_session.query(DocumentChunk) .filter_by(source_id=doc_id, source_type="document") .count() ) # This shows the orphan problem - chunks remain after document deletion assert chunks_after == 3, "Chunks remain orphaned (no FK)" def test_manual_chunk_deletion( self, test_session, test_document, test_collection ): """Verify manual chunk deletion works.""" doc_id = test_document.id collection_name = f"collection_{test_collection.id}" # Create a chunk chunk = DocumentChunk( chunk_hash=f"chunk_hash_{uuid.uuid4().hex[:8]}", source_type="document", source_id=doc_id, collection_name=collection_name, chunk_text="Test chunk content", chunk_index=0, start_char=0, end_char=100, word_count=10, embedding_id=str(uuid.uuid4()), embedding_model="test-model", embedding_model_type="sentence_transformers", ) test_session.add(chunk) test_session.commit() # Manually delete chunks before document deleted = ( test_session.query(DocumentChunk) .filter( DocumentChunk.source_id == doc_id, DocumentChunk.source_type == "document", ) .delete(synchronize_session=False) ) assert deleted == 1, "Should delete 1 chunk" # Now delete document doc = test_session.query(Document).get(doc_id) test_session.delete(doc) test_session.commit() # Verify all cleaned up doc_after = test_session.query(Document).get(doc_id) assert doc_after is None