Cognitive-rag / 01_data_ingestion / tests / test_engine.py
test_engine.py
Raw
"""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"