Cognitive-rag / 01_data_ingestion / tests / conftest.py
conftest.py
Raw
"""Shared fixtures for ingestion tests."""

from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any, Dict, List
from unittest.mock import MagicMock

import pytest

# Ensure the ingestion package is importable
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))


# ---------------------------------------------------------------------------
# Temporary directory fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def tmp_input_dir(tmp_path: Path) -> Path:
    """Create a temp directory with a mix of supported and unsupported files."""
    d = tmp_path / "input"
    d.mkdir()
    return d


@pytest.fixture
def tmp_output_dir(tmp_path: Path) -> Path:
    d = tmp_path / "output"
    d.mkdir()
    return d


@pytest.fixture
def populated_input_dir(tmp_input_dir: Path) -> Path:
    """Input dir with sample files of various types."""
    # Text files
    (tmp_input_dir / "readme.txt").write_text("Hello, this is a test document.")
    (tmp_input_dir / "notes.md").write_text("# Notes\n\nSome markdown content.")

    # CSV file
    (tmp_input_dir / "data.csv").write_text("name,value\nAlice,42\nBob,99\n")

    # Nested subdirectory
    sub = tmp_input_dir / "subdir"
    sub.mkdir()
    (sub / "nested.txt").write_text("Nested content in a subdirectory.")

    # Unsupported files
    (tmp_input_dir / "binary.exe").write_bytes(b"\x00\x01\x02\x03")
    (tmp_input_dir / "archive.tar.gz").write_bytes(b"\x1f\x8b\x08")

    return tmp_input_dir


# ---------------------------------------------------------------------------
# Mock fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def mock_whisper_transcriber():
    """Mock WhisperTranscriber that returns fixed segments."""
    mock = MagicMock()

    # Simulate a segment object
    segment = MagicMock()
    segment.text = "This is a test transcription."
    segment.start = 0.0
    segment.end = 5.0

    info = MagicMock()
    info.language = "en"
    info.language_probability = 0.98

    mock.transcribe.return_value = ([segment], info)
    return mock


@pytest.fixture
def mock_ocr_runner():
    """Mock OcrRunner that returns fixed text."""
    mock = MagicMock()
    mock.engine = "tesseract"
    mock.language = "eng"
    mock.extract_text.return_value = "Sample OCR text from image."
    return mock


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def read_jsonl(path: Path) -> List[Dict[str, Any]]:
    """Read a JSONL file and return list of dicts."""
    records = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records