"""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"