Spaces:
Running on Zero
Running on Zero
File size: 13,300 Bytes
0828c2c | 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 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 | """Tests for document processor module."""
import tempfile
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from langchain_core.documents import Document
from src.document_processor import DocumentProcessor, process_documents
@pytest.fixture
def mock_config():
"""Create a mock config."""
with patch("src.document_processor.get_config") as mock:
mock.return_value.get.side_effect = lambda key, default=None: {
"document_processing.chunk_size": 1000,
"document_processing.chunk_overlap": 200,
"document_processing.supported_extensions": [
".pdf",
".docx",
".doc",
".html",
".htm",
".txt",
".md",
],
"main_document.path": "",
}.get(key, default)
yield mock
@pytest.fixture
def processor(mock_config):
"""Create a DocumentProcessor instance."""
return DocumentProcessor()
# --- DocumentProcessor Tests ---
def test_processor_init(processor):
"""Test DocumentProcessor initialization."""
assert processor is not None
assert processor.chunk_size == 1000
assert processor.chunk_overlap == 200
assert len(processor.supported_extensions) > 0
def test_load_pdf(processor):
"""Test loading PDF file."""
with patch("src.document_processor.PyPDFLoader") as mock_loader:
mock_docs = [
Document(page_content="Page 1 content", metadata={"page": 0}),
Document(page_content="Page 2 content", metadata={"page": 1}),
]
mock_loader.return_value.load.return_value = mock_docs
result = processor.load_pdf(Path("/test/document.pdf"))
assert len(result) == 2
assert result[0].page_content == "Page 1 content"
def test_load_pdf_error(processor):
"""Test loading PDF with error."""
with patch("src.document_processor.PyPDFLoader") as mock_loader:
mock_loader.return_value.load.side_effect = Exception("PDF error")
result = processor.load_pdf(Path("/test/document.pdf"))
assert result == []
def test_load_docx(processor):
"""Test loading DOCX file."""
with patch("src.document_processor.DocxDocument") as mock_docx:
mock_doc = MagicMock()
mock_paragraph1 = MagicMock()
mock_paragraph1.text = "First paragraph"
mock_paragraph2 = MagicMock()
mock_paragraph2.text = "Second paragraph"
mock_doc.paragraphs = [mock_paragraph1, mock_paragraph2]
mock_doc.tables = []
mock_docx.return_value = mock_doc
result = processor.load_docx(Path("/test/document.docx"))
assert len(result) == 1
assert "First paragraph" in result[0].page_content
assert "Second paragraph" in result[0].page_content
def test_load_docx_with_tables(processor):
"""Test loading DOCX file with tables."""
with patch("src.document_processor.DocxDocument") as mock_docx:
mock_doc = MagicMock()
mock_paragraph = MagicMock()
mock_paragraph.text = "Document content"
mock_doc.paragraphs = [mock_paragraph]
# Create mock table
mock_cell1 = MagicMock()
mock_cell1.text = "Cell 1"
mock_cell2 = MagicMock()
mock_cell2.text = "Cell 2"
mock_row = MagicMock()
mock_row.cells = [mock_cell1, mock_cell2]
mock_table = MagicMock()
mock_table.rows = [mock_row]
mock_doc.tables = [mock_table]
mock_docx.return_value = mock_doc
result = processor.load_docx(Path("/test/document.docx"))
assert len(result) == 1
assert "Cell 1" in result[0].page_content
assert "Cell 2" in result[0].page_content
def test_load_docx_error(processor):
"""Test loading DOCX with error."""
with patch("src.document_processor.DocxDocument") as mock_docx:
mock_docx.side_effect = Exception("DOCX error")
result = processor.load_docx(Path("/test/document.docx"))
assert result == []
def test_load_html(processor):
"""Test loading HTML file."""
html_content = """
<html>
<head><style>body { color: red; }</style></head>
<body>
<script>console.log("test");</script>
<h1>Title</h1>
<p>Content paragraph</p>
</body>
</html>
"""
with tempfile.NamedTemporaryFile(mode="w", suffix=".html", delete=False) as f:
f.write(html_content)
f.flush()
result = processor.load_html(Path(f.name))
assert len(result) == 1
assert "Title" in result[0].page_content
assert "Content paragraph" in result[0].page_content
# Script content should be removed
assert "console.log" not in result[0].page_content
def test_load_html_error(processor):
"""Test loading HTML with error."""
result = processor.load_html(Path("/nonexistent/file.html"))
assert result == []
def test_load_text(processor):
"""Test loading text file."""
with patch("src.document_processor.TextLoader") as mock_loader:
mock_docs = [Document(page_content="Text content", metadata={"source": "test.txt"})]
mock_loader.return_value.load.return_value = mock_docs
result = processor.load_text(Path("/test/document.txt"))
assert len(result) == 1
assert result[0].page_content == "Text content"
def test_load_text_error(processor):
"""Test loading text with error."""
with patch("src.document_processor.TextLoader") as mock_loader:
mock_loader.return_value.load.side_effect = Exception("Text error")
result = processor.load_text(Path("/test/document.txt"))
assert result == []
def test_load_document_pdf(processor):
"""Test load_document routes to correct loader for PDF."""
with patch.object(processor, "load_pdf") as mock_load:
mock_load.return_value = [Document(page_content="PDF")]
result = processor.load_document(Path("/test/file.pdf"))
mock_load.assert_called_once()
assert len(result) == 1
def test_load_document_docx(processor):
"""Test load_document routes to correct loader for DOCX."""
with patch.object(processor, "load_docx") as mock_load:
mock_load.return_value = [Document(page_content="DOCX")]
result = processor.load_document(Path("/test/file.docx"))
mock_load.assert_called_once()
assert len(result) == 1
def test_load_document_doc(processor):
"""Test load_document routes to correct loader for DOC."""
with patch.object(processor, "load_docx") as mock_load:
mock_load.return_value = [Document(page_content="DOC")]
processor.load_document(Path("/test/file.doc"))
mock_load.assert_called_once()
def test_load_document_html(processor):
"""Test load_document routes to correct loader for HTML."""
with patch.object(processor, "load_html") as mock_load:
mock_load.return_value = [Document(page_content="HTML")]
processor.load_document(Path("/test/file.html"))
mock_load.assert_called_once()
def test_load_document_htm(processor):
"""Test load_document routes to correct loader for HTM."""
with patch.object(processor, "load_html") as mock_load:
mock_load.return_value = [Document(page_content="HTM")]
processor.load_document(Path("/test/file.htm"))
mock_load.assert_called_once()
def test_load_document_txt(processor):
"""Test load_document routes to correct loader for TXT."""
with patch.object(processor, "load_text") as mock_load:
mock_load.return_value = [Document(page_content="TXT")]
processor.load_document(Path("/test/file.txt"))
mock_load.assert_called_once()
def test_load_document_md(processor):
"""Test load_document routes to correct loader for MD."""
with patch.object(processor, "load_text") as mock_load:
mock_load.return_value = [Document(page_content="MD")]
processor.load_document(Path("/test/file.md"))
mock_load.assert_called_once()
def test_load_document_unsupported(processor):
"""Test load_document returns empty for unsupported extensions."""
result = processor.load_document(Path("/test/file.xyz"))
assert result == []
def test_chunk_documents_empty(processor):
"""Test chunking empty document list."""
result = processor.chunk_documents([])
assert result == []
def test_chunk_documents(processor):
"""Test chunking documents."""
documents = [
Document(
page_content="A" * 2000, # Long content that will be split
metadata={"source": "test.txt"},
)
]
result = processor.chunk_documents(documents)
assert len(result) > 0 # Should be split into multiple chunks
def test_chunk_documents_error(processor):
"""Test chunking with error."""
with patch.object(processor.text_splitter, "split_documents") as mock_split:
mock_split.side_effect = Exception("Chunk error")
result = processor.chunk_documents([Document(page_content="test")])
assert result == []
def test_process_directory_not_found(processor):
"""Test processing non-existent directory."""
result = processor.process_directory("/nonexistent/directory")
assert result == []
def test_process_directory_empty(processor):
"""Test processing empty directory."""
with tempfile.TemporaryDirectory() as tmpdir:
result = processor.process_directory(tmpdir)
assert result == []
def test_process_directory_with_files(processor):
"""Test processing directory with files."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create a test markdown file
test_file = Path(tmpdir) / "test.md"
test_file.write_text("# Test Document\n\nThis is test content.")
with patch.object(processor, "load_text") as mock_load:
mock_load.return_value = [
Document(
page_content="# Test Document\n\nThis is test content.",
metadata={"source": str(test_file)},
)
]
processor.process_directory(tmpdir)
# Should process and chunk the file
mock_load.assert_called()
def test_process_directory_excludes_main_doc():
"""Test processing directory excludes main document."""
with tempfile.TemporaryDirectory() as tmpdir:
# Create test files
main_doc = Path(tmpdir) / "main_profile.md"
main_doc.write_text("# Main Profile")
other_doc = Path(tmpdir) / "other.md"
other_doc.write_text("# Other Document")
with patch("src.document_processor.get_config") as mock_config:
mock_config.return_value.get.side_effect = lambda key, default=None: {
"document_processing.chunk_size": 1000,
"document_processing.chunk_overlap": 200,
"document_processing.supported_extensions": [".md"],
"main_document.path": str(main_doc),
}.get(key, default)
processor = DocumentProcessor()
with patch.object(processor, "load_text") as mock_load:
mock_load.return_value = [Document(page_content="Other content", metadata={})]
processor.process_directory(tmpdir)
# Should only load other.md, not main_profile.md
# Verify main_doc was skipped by checking call args
for call in mock_load.call_args_list:
assert str(main_doc) not in str(call)
# --- process_documents Function Tests ---
def test_process_documents_default_path():
"""Test process_documents uses default path from config."""
with patch("src.document_processor.get_config") as mock_config:
mock_config.return_value.get.side_effect = lambda key, default=None: {
"document_processing.chunk_size": 1000,
"document_processing.chunk_overlap": 200,
"document_processing.supported_extensions": [".md"],
"main_document.path": "",
}.get(key, default)
mock_config.return_value.get_env.return_value = "/default/path"
with patch("src.document_processor.DocumentProcessor") as mock_processor:
mock_processor.return_value.process_directory.return_value = []
process_documents()
mock_processor.return_value.process_directory.assert_called_once_with("/default/path")
def test_process_documents_custom_path():
"""Test process_documents uses custom path."""
with patch("src.document_processor.get_config") as mock_config:
mock_config.return_value.get.side_effect = lambda key, default=None: {
"document_processing.chunk_size": 1000,
"document_processing.chunk_overlap": 200,
"document_processing.supported_extensions": [".md"],
"main_document.path": "",
}.get(key, default)
with patch("src.document_processor.DocumentProcessor") as mock_processor:
mock_processor.return_value.process_directory.return_value = []
process_documents("/custom/path")
mock_processor.return_value.process_directory.assert_called_once_with("/custom/path")
|