Spaces:
Sleeping
Sleeping
File size: 5,417 Bytes
194aaa5 3dbd619 194aaa5 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """Unit tests for DocumentProcessor with mocked dependencies."""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), ".."))
import uuid
import hashlib
import tempfile
from unittest.mock import AsyncMock, MagicMock
import pytest
from models.db import Document
from services.document_processor import DocumentProcessor
@pytest.fixture
def mock_db():
"""Create a mocked async database session."""
db = AsyncMock()
# Make db.add a no-op side effect that actually stores the object
db.add = MagicMock()
# Default execute to return a result that yields None for scalars().first()
mock_result = MagicMock()
mock_result.scalars.return_value.first.return_value = None
db.execute.return_value = mock_result
return db
@pytest.fixture
def mock_embedding_service():
"""Create a mocked embedding service that returns fake embeddings."""
service = MagicMock()
service.embed_chunks.return_value = [[0.1] * 384, [0.2] * 384]
return service
@pytest.fixture
def sample_txt():
"""Create a temporary text file for testing."""
with tempfile.NamedTemporaryFile(
mode="w", suffix=".txt", delete=False, encoding="utf-8"
) as f:
f.write("This is a test document with enough words to create multiple chunks for testing the document processor pipeline.")
f.write(" " + "word " * 500) # Add more words to ensure chunking
tmp_path = f.name
yield tmp_path
if os.path.exists(tmp_path):
os.unlink(tmp_path)
@pytest.mark.asyncio
async def test_process_document_success(mock_db, mock_embedding_service, sample_txt):
"""Should successfully process a txt file through the full pipeline."""
user_id = uuid.uuid4()
processor = DocumentProcessor(mock_db, mock_embedding_service, user_id)
# Configure mock flush to set document.id
async def flush_side_effect():
pass
mock_db.flush = AsyncMock(side_effect=flush_side_effect)
doc = await processor.process_document(sample_txt, "test.txt")
assert doc.filename == "test.txt"
assert doc.status == "ready"
assert doc.content_hash is not None
assert doc.user_id == user_id
# Verify that chunks were added
chunk_add_calls = [call for call in mock_db.add.call_args_list if call[0][0].__class__.__name__ == "Chunk"]
assert len(chunk_add_calls) > 0
mock_db.commit.assert_awaited_once()
@pytest.mark.asyncio
async def test_process_document_empty_file(mock_db, mock_embedding_service):
"""Should raise ValueError on empty document."""
with tempfile.NamedTemporaryFile(
mode="w", suffix=".txt", delete=False, encoding="utf-8"
) as f:
f.write(" ")
empty_path = f.name
user_id = uuid.uuid4()
processor = DocumentProcessor(mock_db, mock_embedding_service, user_id)
try:
with pytest.raises(ValueError, match="empty"):
await processor.process_document(empty_path, "empty.txt")
finally:
os.unlink(empty_path)
@pytest.mark.asyncio
async def test_process_document_unsupported_extension(mock_db, mock_embedding_service):
"""Should raise ValueError for unsupported file extensions."""
with tempfile.NamedTemporaryFile(
mode="w", suffix=".xyz", delete=False
) as f:
f.write("some text")
bad_path = f.name
user_id = uuid.uuid4()
processor = DocumentProcessor(mock_db, mock_embedding_service, user_id)
try:
with pytest.raises(ValueError, match="Unsupported"):
await processor.process_document(bad_path, "test.xyz")
finally:
os.unlink(bad_path)
@pytest.mark.asyncio
async def test_process_document_dedup(mock_db, mock_embedding_service, sample_txt):
"""Should return existing document if content hash matches."""
user_id = uuid.uuid4()
processor = DocumentProcessor(mock_db, mock_embedding_service, user_id)
with open(sample_txt, "r") as f:
text = f.read()
content_hash = hashlib.sha256(text.encode("utf-8")).hexdigest()
# Create an "existing" document that the DB will return
existing_doc = Document(
id=uuid.uuid4(),
user_id=user_id,
filename="existing.txt",
content_hash=content_hash,
file_path=sample_txt,
status="ready",
)
# Mock the DB to return existing_doc on the first select query
mock_result = MagicMock()
mock_result.scalars.return_value.first.return_value = existing_doc
mock_db.execute.return_value = mock_result
result = await processor.process_document(sample_txt, "test.txt")
# Should return the existing document without creating a new one
assert result.id == existing_doc.id
assert result.filename == "existing.txt"
# The select query should have been executed
mock_db.execute.assert_awaited_once()
@pytest.mark.asyncio
async def test_process_document_user_id_none(mock_db, mock_embedding_service, sample_txt):
"""Should handle user_id=None gracefully (guest users)."""
processor = DocumentProcessor(mock_db, mock_embedding_service, user_id=None)
async def flush_side_effect():
pass
mock_db.flush = AsyncMock(side_effect=flush_side_effect)
doc = await processor.process_document(sample_txt, "guest_test.txt")
assert doc.filename == "guest_test.txt"
assert doc.status == "ready"
mock_db.commit.assert_awaited_once()
|