"""Tests for the parallel IngestionEngine.""" from __future__ import annotations import json import sys import time from pathlib import Path from typing import Any, Dict, List from unittest.mock import MagicMock, patch import pytest sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from ingestion.engine import IngestionConfig, IngestionEngine, FileResult from ingestion.events import ( CallbackEmitter, FileCompleted, FileEmpty, FileFailed, FileSkipped, FileStarted, IngestionCompleted, IngestionProgress, IngestionStarted, NullEmitter, ) from tests.conftest import read_jsonl # --------------------------------------------------------------------------- # Fixtures # --------------------------------------------------------------------------- @pytest.fixture def basic_config(tmp_path: Path) -> IngestionConfig: input_dir = tmp_path / "input" input_dir.mkdir() output_dir = tmp_path / "output" output_dir.mkdir() return IngestionConfig(input_dir=input_dir, output_dir=output_dir, workers=2) @pytest.fixture def txt_csv_input(basic_config: IngestionConfig) -> IngestionConfig: """Add text and CSV files to the input directory.""" d = basic_config.input_dir (d / "hello.txt").write_text("Hello world from a test file.") (d / "data.csv").write_text("name,value\nAlice,42\nBob,99\n") return basic_config @pytest.fixture def mixed_input(basic_config: IngestionConfig) -> IngestionConfig: """Add supported + unsupported files.""" d = basic_config.input_dir (d / "doc.txt").write_text("Some text content.") (d / "data.csv").write_text("col1,col2\na,b\n") (d / "binary.exe").write_bytes(b"\x00\x01") (d / "archive.tar.gz").write_bytes(b"\x1f\x8b") sub = d / "nested" sub.mkdir() (sub / "deep.txt").write_text("Nested file content.") return basic_config # --------------------------------------------------------------------------- # File classification # --------------------------------------------------------------------------- class TestFileClassification: def test_supported_and_skipped(self, mixed_input: IngestionConfig): engine = IngestionEngine(mixed_input) from ingest_multimodal import iter_files files = list(iter_files(mixed_input.input_dir)) supported, skipped = engine._classify_files(files) supported_exts = {Path(s[1]).suffix for s in supported} skipped_exts = {s["path"].rsplit(".", 1)[-1] for s in skipped} assert ".txt" in supported_exts assert ".csv" in supported_exts assert "exe" in skipped_exts or "tar" in skipped_exts def test_nested_files_discovered(self, mixed_input: IngestionConfig): engine = IngestionEngine(mixed_input) from ingest_multimodal import iter_files files = list(iter_files(mixed_input.input_dir)) supported, _ = engine._classify_files(files) paths = [str(s[1]) for s in supported] assert any("nested" in p and "deep.txt" in p for p in paths) # --------------------------------------------------------------------------- # Sequential vs parallel output consistency # --------------------------------------------------------------------------- class TestOutputConsistency: def test_sequential_produces_output(self, txt_csv_input: IngestionConfig): txt_csv_input.workers = 0 engine = IngestionEngine(txt_csv_input) stats = engine.run() assert stats["counts"]["files_succeeded"] >= 1 assert stats["counts"]["records_emitted"] >= 1 output_path = txt_csv_input.output_dir / "cleaned_documents.jsonl" assert output_path.exists() records = read_jsonl(output_path) assert len(records) > 0 def test_parallel_produces_output(self, txt_csv_input: IngestionConfig): txt_csv_input.workers = 2 engine = IngestionEngine(txt_csv_input) stats = engine.run() assert stats["counts"]["files_succeeded"] >= 1 assert stats["counts"]["records_emitted"] >= 1 def test_sequential_and_parallel_same_records(self, tmp_path: Path): """Both modes should produce the same JSONL content.""" # Create shared input input_dir = tmp_path / "input" input_dir.mkdir() (input_dir / "a.txt").write_text("Content A") (input_dir / "b.csv").write_text("x,y\n1,2\n3,4\n") (input_dir / "c.txt").write_text("Content C") # Sequential out_seq = tmp_path / "out_seq" out_seq.mkdir() cfg_seq = IngestionConfig(input_dir=input_dir, output_dir=out_seq, workers=0) IngestionEngine(cfg_seq).run() # Parallel out_par = tmp_path / "out_par" out_par.mkdir() cfg_par = IngestionConfig(input_dir=input_dir, output_dir=out_par, workers=4) IngestionEngine(cfg_par).run() recs_seq = read_jsonl(out_seq / "cleaned_documents.jsonl") recs_par = read_jsonl(out_par / "cleaned_documents.jsonl") # Same number of records assert len(recs_seq) == len(recs_par) # Same doc_ids and content (order should match since we sort by file index) seq_ids = [r["doc_id"] for r in recs_seq] par_ids = [r["doc_id"] for r in recs_par] assert seq_ids == par_ids # --------------------------------------------------------------------------- # Event emission # --------------------------------------------------------------------------- class TestEventEmission: def test_emits_lifecycle_events(self, txt_csv_input: IngestionConfig): received = [] emitter = CallbackEmitter(received.append) engine = IngestionEngine(txt_csv_input, emitter=emitter) engine.run() types = [e.event_type for e in received] assert "ingestion_started" in types assert "file_started" in types assert "file_completed" in types assert "ingestion_progress" in types assert "ingestion_completed" in types def test_skipped_files_emit_events(self, mixed_input: IngestionConfig): received = [] emitter = CallbackEmitter(received.append) engine = IngestionEngine(mixed_input, emitter=emitter) engine.run() skip_events = [e for e in received if e.event_type == "file_skipped"] assert len(skip_events) >= 1 def test_null_emitter_works(self, txt_csv_input: IngestionConfig): engine = IngestionEngine(txt_csv_input, emitter=NullEmitter()) stats = engine.run() assert stats["counts"]["records_emitted"] >= 1 # --------------------------------------------------------------------------- # Error handling # --------------------------------------------------------------------------- class TestErrorHandling: def test_error_in_one_file_does_not_stop_others(self, basic_config: IngestionConfig): d = basic_config.input_dir (d / "good.csv").write_text("a,b\n1,2\n") # Create a CSV that will fail (binary content as CSV) (d / "bad.csv").write_bytes(b"\x00\x01\x02" * 100) (d / "also_good.txt").write_text("This is fine.") engine = IngestionEngine(basic_config) stats = engine.run() # At least one file should succeed even if another fails assert stats["counts"]["files_succeeded"] >= 1 def test_empty_input_dir(self, basic_config: IngestionConfig): engine = IngestionEngine(basic_config) stats = engine.run() assert stats["counts"]["files_processed"] == 0 assert stats["counts"]["records_emitted"] == 0 # --------------------------------------------------------------------------- # Cancellation # --------------------------------------------------------------------------- class TestCancellation: def test_cancel_stops_processing(self, basic_config: IngestionConfig): d = basic_config.input_dir for i in range(20): (d / f"file_{i:03d}.txt").write_text(f"Content {i}") basic_config.workers = 1 engine = IngestionEngine(basic_config) # Cancel immediately engine.cancel() stats = engine.run() # Should have processed fewer than all files assert stats["counts"]["files_succeeded"] < 20 def test_is_cancelled_property(self, basic_config: IngestionConfig): engine = IngestionEngine(basic_config) assert not engine.is_cancelled engine.cancel() assert engine.is_cancelled # --------------------------------------------------------------------------- # Stats output # --------------------------------------------------------------------------- class TestStatsOutput: def test_stats_json_written(self, txt_csv_input: IngestionConfig): engine = IngestionEngine(txt_csv_input) engine.run() stats_path = txt_csv_input.output_dir / "ingestion_stats.json" assert stats_path.exists() with open(stats_path) as f: stats = json.load(f) assert "counts" in stats assert "config" in stats assert stats["config"]["workers"] == 2 assert "duration_sec" in stats assert "generated_at" in stats def test_errors_jsonl_written(self, basic_config: IngestionConfig): engine = IngestionEngine(basic_config) engine.run() errors_path = basic_config.output_dir / "ingestion_errors.jsonl" assert errors_path.exists() def test_skipped_files_in_stats(self, mixed_input: IngestionConfig): engine = IngestionEngine(mixed_input) stats = engine.run() assert stats["counts"]["files_skipped"] >= 1 assert len(stats["skipped_files"]) >= 1 assert all("reason" in s for s in stats["skipped_files"]) # --------------------------------------------------------------------------- # Config dataclass # --------------------------------------------------------------------------- class TestIngestionConfig: def test_defaults(self, tmp_path: Path): cfg = IngestionConfig(input_dir=tmp_path, output_dir=tmp_path) assert cfg.workers == 4 assert cfg.pdf_strategy == "auto" assert cfg.whisper_model == "small" assert cfg.whisper_device == "cpu" def test_custom_values(self, tmp_path: Path): cfg = IngestionConfig( input_dir=tmp_path, output_dir=tmp_path, workers=8, whisper_device="cuda", ) assert cfg.workers == 8 assert cfg.whisper_device == "cuda"