#!/usr/bin/env python
"""Automated test suite for the multimodal ingestion pipeline.
Runs per-modality tests against 00_Data_test/ datasets, validates output
schema, and reports clear pass/fail results.
Usage:
python test_ingestion.py --data-dir 00_Data_test --modality all
python test_ingestion.py --modality text --verbose
python test_ingestion.py --modality image
"""
from __future__ import annotations
import argparse
import json
import shutil
import sys
import tempfile
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple
# Ensure sibling module is importable
sys.path.insert(0, str(Path(__file__).resolve().parent))
from ingest_multimodal import ingest, detect_file_type, iter_files
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
REQUIRED_FIELDS = {"doc_id", "page", "section", "section_index", "content", "source"}
REQUIRED_SOURCE_FIELDS = {"path", "filename", "filetype", "modality"}
def read_jsonl(path: Path) -> List[Dict]:
records = []
if not path.exists():
return 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
def validate_record(record: Dict) -> List[str]:
"""Return a list of schema violation messages for a single record."""
issues = []
missing = REQUIRED_FIELDS - set(record)
if missing:
issues.append(f"Missing fields: {sorted(missing)}")
if not record.get("doc_id"):
issues.append("Empty doc_id")
if not record.get("content"):
issues.append("Empty content")
if not record.get("section"):
issues.append("Empty section")
source = record.get("source")
if not isinstance(source, dict):
issues.append("source is not a dict")
else:
for field in REQUIRED_SOURCE_FIELDS:
if not source.get(field):
issues.append(f"Missing source.{field}")
# Video-specific validation
modality = source.get("modality")
if modality in ("video", "video_frame"):
if source.get("timestamp_start") is None:
issues.append("video record missing timestamp_start")
if modality == "video_frame":
if not source.get("frame_path"):
issues.append("frame record missing frame_path")
return issues
def has_faster_whisper() -> bool:
try:
import faster_whisper # noqa: F401
return True
except Exception:
return False
def has_ffmpeg() -> bool:
return shutil.which("ffmpeg") is not None
def has_tesseract() -> bool:
return shutil.which("tesseract") is not None
def run_ingest(
input_dir: Path,
max_files: Optional[int] = None,
whisper_model: str = "tiny",
chunk_duration: float = 600.0,
extract_keyframes: bool = False,
scene_threshold: float = 0.3,
) -> Tuple[Dict, List[Dict], List[Dict], Path]:
"""Run ingestion into a temp directory and return (stats, records, errors, tmp_dir).
Note: returns the tmp_dir path so callers can inspect output files.
The caller is responsible for cleanup if needed (but tempfile handles it).
"""
tmp_dir_obj = tempfile.mkdtemp()
tmp = Path(tmp_dir_obj)
stats = ingest(
input_dir=input_dir,
output_dir=tmp,
output_name="cleaned_documents.jsonl",
errors_name="ingestion_errors.jsonl",
pdf_strategy="auto",
ocr_engine="tesseract",
ocr_language="eng",
whisper_model=whisper_model,
whisper_device="cpu",
whisper_compute_type="int8",
max_files=max_files,
chunk_duration=chunk_duration,
extract_keyframes=extract_keyframes,
scene_threshold=scene_threshold,
)
records = read_jsonl(tmp / "cleaned_documents.jsonl")
errors = read_jsonl(tmp / "ingestion_errors.jsonl")
return stats, records, errors, tmp
# ---------------------------------------------------------------------------
# Individual tests
# ---------------------------------------------------------------------------
def test_text(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test text ingestion (WikiText-103 .txt files)."""
input_dir = data_dir / "text" / "wikitext-103"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --text", None
stats, records, errors, _tmp = run_ingest(input_dir)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
errs = stats["counts"]["errors"]
if errs > 0:
return "FAIL", f"{files} files -> {recs} records, {errs} errors", records
if recs == 0:
return "FAIL", f"{files} files -> 0 records (expected > 0)", records
# Verify modality
for rec in records:
src = rec.get("source", {})
if src.get("modality") != "text":
return "FAIL", f"Expected modality 'text', got '{src.get('modality')}'", records
return "PASS", f"{files} files -> {recs} records, {errs} errors", records
def test_image(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test image OCR ingestion (CORD-v2 receipt images)."""
input_dir = data_dir / "image" / "cord-v2" / "images"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --image", None
if not has_tesseract():
return "SKIP", "tesseract not installed", None
stats, records, errors, _tmp = run_ingest(input_dir, max_files=5)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
n_failed = stats["counts"]["files_failed"]
n_empty = stats["counts"]["files_empty"]
if n_failed > 0:
return "FAIL", f"{files} files -> {recs} records, {n_failed} failures", records
# Some images may produce empty OCR — that's expected and tracked in empty_files
if recs == 0:
return "FAIL", f"{files} files -> 0 records (expected >= 1)", records
for rec in records:
src = rec.get("source", {})
if src.get("modality") != "image":
return "FAIL", f"Expected modality 'image', got '{src.get('modality')}'", records
if not src.get("ocr_engine"):
return "FAIL", "Missing source.ocr_engine on image record", records
detail = f"{files} files -> {recs} records, {n_empty} empty OCR"
return "PASS", detail, records
def test_audio(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test audio transcription (LibriSpeech FLAC files)."""
input_dir = data_dir / "audio_subset"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --audio", None
if not has_faster_whisper():
return "SKIP", "faster-whisper not available (requires Python <3.13)", None
stats, records, errors, _tmp = run_ingest(input_dir)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
errs = stats["counts"]["errors"]
if errs > 0:
return "FAIL", f"{files} files -> {recs} records, {errs} errors", records
if recs == 0:
return "FAIL", f"{files} files -> 0 records (expected > 0)", records
for rec in records:
src = rec.get("source", {})
if src.get("modality") != "audio":
return "FAIL", f"Expected modality 'audio', got '{src.get('modality')}'", records
if src.get("timestamp_start") is None:
return "FAIL", "Missing source.timestamp_start on audio record", records
return "PASS", f"{files} files -> {recs} records, {errs} errors", records
def test_video(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test video transcription (YouTube .webm files)."""
input_dir = data_dir / "video" / "yt-ugc-subset"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --video", None
if not has_faster_whisper():
return "SKIP", "faster-whisper not available (requires Python <3.13)", None
if not has_ffmpeg():
return "SKIP", "ffmpeg not installed", None
stats, records, errors, _tmp = run_ingest(input_dir, max_files=2)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
errs = stats["counts"]["errors"]
if errs > 0:
return "FAIL", f"{files} files -> {recs} records, {errs} errors", records
if recs == 0:
return "FAIL", f"{files} files -> 0 records (expected > 0)", records
for rec in records:
src = rec.get("source", {})
if src.get("modality") != "video":
return "FAIL", f"Expected modality 'video', got '{src.get('modality')}'", records
if src.get("filetype") != "video":
return "FAIL", f"Expected filetype 'video', got '{src.get('filetype')}'", records
return "PASS", f"{files} files -> {recs} records, {errs} errors", records
def test_csv(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test CSV ingestion (mtsamples.csv)."""
# CSV lives in 00_Data, one level up from test data
csv_dir = data_dir.parent / "00_Data"
if not csv_dir.exists() or not (csv_dir / "mtsamples.csv").exists():
return "SKIP", "00_Data/mtsamples.csv not found", None
stats, records, errors, _tmp = run_ingest(csv_dir, max_files=1)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
errs = stats["counts"]["errors"]
if errs > 0:
return "FAIL", f"{files} files -> {recs} records, {errs} errors", records
if recs == 0:
return "FAIL", f"{files} files -> 0 records (expected > 0)", records
for rec in records[:5]: # spot-check first 5
src = rec.get("source", {})
if src.get("modality") != "table":
return "FAIL", f"Expected modality 'table', got '{src.get('modality')}'", records
return "PASS", f"{files} files -> {recs} records, {errs} errors", records
def test_mixed(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test mixed-modality batch processing across 00_Data_test/."""
if not data_dir.exists():
return "SKIP", "Test data directory not found", None
stats, records, errors, _tmp = run_ingest(data_dir, max_files=15)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
errs = stats["counts"]["errors"]
# In mixed mode, some errors are acceptable (e.g. missing faster-whisper)
# but we should get at least some records
if recs == 0:
return "FAIL", f"{files} files -> 0 records, {errs} errors", records
# Check that multiple modalities appear
modalities = {rec.get("source", {}).get("modality") for rec in records}
modalities.discard(None)
detail = f"{files} files -> {recs} records, {errs} errors, modalities: {sorted(modalities)}"
if len(modalities) < 1:
return "FAIL", f"No modalities detected. {detail}", records
return "PASS", detail, records
def test_max_files(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test that max_files cap works correctly at different sizes."""
input_dir = data_dir / "text" / "wikitext-103"
if not input_dir.exists():
return "SKIP", "Test data not found", None
all_records = []
for cap in [1, 2, 3]:
stats, records, _, _tmp = run_ingest(input_dir, max_files=cap)
files = stats["counts"]["files_processed"]
if files > cap:
return "FAIL", f"max_files={cap} but processed {files} files", records
all_records.extend(records)
# Also test None (no cap) — should process all files
stats_all, records_all, _, _tmp = run_ingest(input_dir, max_files=None)
total = stats_all["counts"]["files_processed"]
return "PASS", f"max_files caps 1,2,3 respected; uncapped processed {total} files", all_records
def test_reporting(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test that failure reporting includes success_files, failed_files, empty_files."""
input_dir = data_dir / "image" / "cord-v2" / "images"
if not input_dir.exists():
return "SKIP", "Test data not found", None
if not has_tesseract():
return "SKIP", "tesseract not installed", None
stats, records, errors, _tmp = run_ingest(input_dir, max_files=5)
# Verify new stats fields exist
counts = stats.get("counts", {})
for field in ("files_succeeded", "files_failed", "files_empty"):
if field not in counts:
return "FAIL", f"Missing counts.{field} in stats", records
# Verify file lists exist
for field in ("success_files", "failed_files", "empty_files"):
if field not in stats:
return "FAIL", f"Missing {field} list in stats", records
# Verify counts are consistent
total = counts["files_succeeded"] + counts["files_failed"] + counts["files_empty"]
if total != counts["files_processed"]:
return "FAIL", (f"Counts inconsistent: succeeded({counts['files_succeeded']}) + "
f"failed({counts['files_failed']}) + empty({counts['files_empty']}) "
f"!= processed({counts['files_processed']})"), records
# Verify empty_files entries have modality and reason
for entry in stats["empty_files"]:
if "modality" not in entry or "reason" not in entry:
return "FAIL", f"empty_files entry missing modality/reason: {entry}", records
# Verify failed_files entries have modality and reason (if any)
for entry in stats["failed_files"]:
if "modality" not in entry or "reason" not in entry:
return "FAIL", f"failed_files entry missing modality/reason: {entry}", records
# Verify error log entries include modality
for err in errors:
if "modality" not in err:
return "FAIL", f"Error log entry missing modality: {err}", records
detail = (f"succeeded={counts['files_succeeded']}, "
f"failed={counts['files_failed']}, "
f"empty={counts['files_empty']}")
return "PASS", detail, records
def test_video_chunked(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test that chunked video processing produces correct timestamps."""
input_dir = data_dir / "video" / "yt-ugc-subset"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --video", None
if not has_faster_whisper():
return "SKIP", "faster-whisper not available (requires Python <3.13)", None
if not has_ffmpeg():
return "SKIP", "ffmpeg not installed", None
# Use a very short chunk duration to force multiple chunks
stats, records, errors, tmp = run_ingest(
input_dir, max_files=2, chunk_duration=30.0,
)
files = stats["counts"]["files_processed"]
recs = stats["counts"]["records_emitted"]
if recs == 0:
return "FAIL", f"{files} files -> 0 records (expected > 0)", records
# Verify timestamps are monotonically non-decreasing per doc_id
by_doc: Dict[str, List[float]] = {}
for rec in records:
src = rec.get("source", {})
if src.get("modality") != "video":
continue
doc = rec["doc_id"]
by_doc.setdefault(doc, []).append(src.get("timestamp_start", 0.0))
for doc, timestamps in by_doc.items():
for i in range(1, len(timestamps)):
if timestamps[i] < timestamps[i - 1]:
return "FAIL", f"Timestamps not monotonic for {doc} at index {i}", records
# Verify segment numbering is continuous per doc_id
seg_by_doc: Dict[str, List[str]] = {}
for rec in records:
if rec.get("source", {}).get("modality") == "video":
seg_by_doc.setdefault(rec["doc_id"], []).append(rec["section"])
for doc, sections in seg_by_doc.items():
expected = [f"segment_{i:04d}" for i in range(len(sections))]
if sections != expected:
return "FAIL", f"Segment numbering not continuous for {doc}", records
# Verify no checkpoint files remain after success
checkpoints = list(tmp.glob(".checkpoint_*.json"))
if checkpoints:
return "FAIL", f"Checkpoint files not cleaned up: {checkpoints}", records
# Clean up
shutil.rmtree(tmp, ignore_errors=True)
return "PASS", f"{recs} records with monotonic timestamps across chunks", records
def test_video_resume(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test that checkpoint resume works correctly."""
input_dir = data_dir / "video" / "yt-ugc-subset"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --video", None
if not has_faster_whisper():
return "SKIP", "faster-whisper not available (requires Python <3.13)", None
if not has_ffmpeg():
return "SKIP", "ffmpeg not installed", None
# Run a full fresh ingestion with short chunks
_stats1, records1, _, tmp1 = run_ingest(
input_dir, max_files=1, chunk_duration=30.0,
)
shutil.rmtree(tmp1, ignore_errors=True)
if not records1:
return "SKIP", "No records produced (video may be too short)", None
# Run again — should produce same number of records
_stats2, records2, _, tmp2 = run_ingest(
input_dir, max_files=1, chunk_duration=30.0,
)
shutil.rmtree(tmp2, ignore_errors=True)
if len(records1) != len(records2):
return "FAIL", (f"Inconsistent record counts: "
f"run1={len(records1)}, run2={len(records2)}"), records2
return "PASS", f"Consistent output across runs ({len(records1)} records)", records1
def test_keyframes(data_dir: Path, verbose: bool) -> Tuple[str, str, Optional[List[Dict]]]:
"""Test keyframe extraction with OCR produces frame records."""
input_dir = data_dir / "video" / "yt-ugc-subset"
if not input_dir.exists():
return "SKIP", "Test data not found. Run download_test_datasets.py --video", None
if not has_faster_whisper():
return "SKIP", "faster-whisper not available (requires Python <3.13)", None
if not has_ffmpeg():
return "SKIP", "ffmpeg not installed", None
if not has_tesseract():
return "SKIP", "tesseract not installed", None
_stats, records, _errors, tmp = run_ingest(
input_dir, max_files=1, extract_keyframes=True, scene_threshold=0.3,
)
segment_records = [r for r in records if r["section"].startswith("segment_")]
frame_records = [r for r in records if r["section"].startswith("frame_")]
if not segment_records:
shutil.rmtree(tmp, ignore_errors=True)
return "FAIL", "No transcript segments produced", records
# Verify frame record schema if any frames were extracted
for rec in frame_records:
src = rec.get("source", {})
if src.get("modality") != "video_frame":
shutil.rmtree(tmp, ignore_errors=True)
return "FAIL", f"Frame record has wrong modality: {src.get('modality')}", records
if not src.get("frame_path"):
shutil.rmtree(tmp, ignore_errors=True)
return "FAIL", "Frame record missing frame_path", records
if src.get("timestamp_start") is None:
shutil.rmtree(tmp, ignore_errors=True)
return "FAIL", "Frame record missing timestamp", records
if not src.get("ocr_engine"):
shutil.rmtree(tmp, ignore_errors=True)
return "FAIL", "Frame record missing ocr_engine", records
# Check that frames directory was created if frames were extracted
if frame_records:
frames_dirs = list(tmp.glob("frames/*"))
if not frames_dirs:
shutil.rmtree(tmp, ignore_errors=True)
return "FAIL", "Frame records exist but no frames/ directory found", records
shutil.rmtree(tmp, ignore_errors=True)
detail = f"{len(segment_records)} segments, {len(frame_records)} frames"
return "PASS", detail, records
def test_schema(all_records: List[Dict], verbose: bool) -> Tuple[str, str, None]:
"""Validate schema compliance across all collected records."""
if not all_records:
return "SKIP", "No records to validate", None
total_issues = 0
issue_details = []
for i, rec in enumerate(all_records):
issues = validate_record(rec)
if issues:
total_issues += len(issues)
if verbose and len(issue_details) < 10:
issue_details.append(f" Record {i}: {'; '.join(issues)}")
if total_issues > 0:
detail = f"{total_issues} issues across {len(all_records)} records"
if issue_details:
detail += "\n" + "\n".join(issue_details)
return "FAIL", detail, None
return "PASS", f"All {len(all_records)} records pass schema validation", None
# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------
TESTS = {
"text": test_text,
"image": test_image,
"audio": test_audio,
"video": test_video,
"csv": test_csv,
"mixed": test_mixed,
"max_files": test_max_files,
"reporting": test_reporting,
"video_chunked": test_video_chunked,
"video_resume": test_video_resume,
"keyframes": test_keyframes,
}
def run_tests(data_dir: Path, modality: str, verbose: bool) -> int:
"""Run selected tests and return exit code (0=pass, 1=failures)."""
print("=== Multimodal Ingestion Test Suite ===\n")
if modality == "all":
test_names = list(TESTS.keys())
else:
test_names = [modality]
passed = 0
failed = 0
skipped = 0
all_records: List[Dict] = []
for name in test_names:
test_fn = TESTS.get(name)
if test_fn is None:
print(f"[????] Unknown test: {name}")
continue
start = time.time()
try:
status, message, records = test_fn(data_dir, verbose)
except Exception as exc:
status = "FAIL"
message = f"Unhandled exception: {exc}"
records = None
elapsed = time.time() - start
if records:
all_records.extend(records)
tag = {"PASS": "PASS", "FAIL": "FAIL", "SKIP": "SKIP"}.get(status, "????")
time_str = f"({elapsed:.1f}s)" if status != "SKIP" else ""
print(f"[{tag}] {name:<12} {message} {time_str}")
if verbose and status == "FAIL" and records:
for rec in records[:3]:
print(f" Sample record: doc_id={rec.get('doc_id')}, "
f"modality={rec.get('source', {}).get('modality')}")
if status == "PASS":
passed += 1
elif status == "FAIL":
failed += 1
else:
skipped += 1
# Schema test runs on all collected records
if modality == "all" and all_records:
start = time.time()
status, message, _ = test_schema(all_records, verbose)
elapsed = time.time() - start
tag = {"PASS": "PASS", "FAIL": "FAIL", "SKIP": "SKIP"}.get(status, "????")
print(f"[{tag}] {'schema':<12} {message} ({elapsed:.1f}s)")
if status == "PASS":
passed += 1
elif status == "FAIL":
failed += 1
else:
skipped += 1
print(f"\nResults: {passed} passed, {failed} failed, {skipped} skipped")
return 1 if failed > 0 else 0
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Test the multimodal ingestion pipeline.")
parser.add_argument(
"--data-dir",
default="00_Data_test",
help="Root directory containing test data (default: 00_Data_test).",
)
parser.add_argument(
"--modality",
choices=["text", "image", "audio", "video", "csv", "mixed", "max_files",
"reporting", "video_chunked", "video_resume", "keyframes", "all"],
default="all",
help="Which test(s) to run (default: all).",
)
parser.add_argument(
"--verbose",
action="store_true",
help="Print detailed output for each test.",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
data_dir = Path(args.data_dir)
if not data_dir.exists():
print(f"Error: data directory '{data_dir}' not found.")
print("Run download_test_datasets.py first to download test data.")
sys.exit(1)
exit_code = run_tests(data_dir, args.modality, args.verbose)
sys.exit(exit_code)
if __name__ == "__main__":
main()