Cognitive-rag / 01_data_ingestion / tests / test_api.py
test_api.py
Raw
"""Tests for the ingestion API endpoints."""

from __future__ import annotations

import io
import json
import sys
import time
import zipfile
from pathlib import Path
from unittest.mock import patch

import pytest

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from fastapi import FastAPI
from fastapi.testclient import TestClient

from ingestion.api import router, _jobs


@pytest.fixture(autouse=True)
def clear_jobs():
    """Clear the job store between tests."""
    _jobs.clear()
    yield
    _jobs.clear()


@pytest.fixture
def app():
    app = FastAPI()
    app.include_router(router)
    return app


@pytest.fixture
def client(app):
    return TestClient(app)


@pytest.fixture
def input_dir_with_files(tmp_path: Path) -> Path:
    d = tmp_path / "test_input"
    d.mkdir()
    (d / "hello.txt").write_text("Hello world")
    (d / "data.csv").write_text("a,b\n1,2\n")
    return d


# ---------------------------------------------------------------------------
# POST /ingestion/start
# ---------------------------------------------------------------------------

class TestStartIngestion:
    def test_start_valid_path(self, client, input_dir_with_files):
        resp = client.post("/ingestion/start", json={
            "input_dir": str(input_dir_with_files),
            "output_dir": str(input_dir_with_files.parent / "output"),
            "workers": 2,
        })
        assert resp.status_code == 200
        data = resp.json()
        assert "job_id" in data
        assert data["status"] in ("pending", "running")

    def test_start_invalid_path(self, client):
        resp = client.post("/ingestion/start", json={
            "input_dir": "/nonexistent/path/that/does/not/exist",
        })
        assert resp.status_code == 400
        assert "not found" in resp.json()["detail"].lower()

    def test_start_returns_unique_job_ids(self, client, input_dir_with_files):
        r1 = client.post("/ingestion/start", json={"input_dir": str(input_dir_with_files)})
        r2 = client.post("/ingestion/start", json={"input_dir": str(input_dir_with_files)})
        assert r1.json()["job_id"] != r2.json()["job_id"]


# ---------------------------------------------------------------------------
# POST /ingestion/upload
# ---------------------------------------------------------------------------

class TestUploadIngestion:
    def test_upload_single_file(self, client, tmp_path: Path):
        output_dir = tmp_path / "upload_output"
        output_dir.mkdir()
        resp = client.post(
            "/ingestion/upload",
            files=[("files", ("test.txt", b"Hello from upload", "text/plain"))],
            data={"output_dir": str(output_dir), "workers": "1"},
        )
        assert resp.status_code == 200
        data = resp.json()
        assert data["uploaded_files"] >= 1
        assert "job_id" in data

    def test_upload_zip_file(self, client, tmp_path: Path):
        output_dir = tmp_path / "zip_output"
        output_dir.mkdir()

        # Create a zip in memory
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w") as zf:
            zf.writestr("doc1.txt", "First document")
            zf.writestr("subdir/doc2.txt", "Second document")
        buf.seek(0)

        resp = client.post(
            "/ingestion/upload",
            files=[("files", ("archive.zip", buf.getvalue(), "application/zip"))],
            data={"output_dir": str(output_dir), "workers": "1"},
        )
        assert resp.status_code == 200
        data = resp.json()
        assert data["uploaded_files"] >= 2

    def test_upload_multiple_files(self, client, tmp_path: Path):
        output_dir = tmp_path / "multi_output"
        output_dir.mkdir()
        resp = client.post(
            "/ingestion/upload",
            files=[
                ("files", ("a.txt", b"File A content", "text/plain")),
                ("files", ("b.csv", b"x,y\n1,2\n", "text/csv")),
            ],
            data={"output_dir": str(output_dir), "workers": "1"},
        )
        assert resp.status_code == 200
        assert resp.json()["uploaded_files"] == 2


# ---------------------------------------------------------------------------
# GET /ingestion/{job_id}/status
# ---------------------------------------------------------------------------

class TestJobStatus:
    def test_status_of_running_job(self, client, input_dir_with_files):
        start_resp = client.post("/ingestion/start", json={
            "input_dir": str(input_dir_with_files),
            "output_dir": str(input_dir_with_files.parent / "output"),
            "workers": 1,
        })
        job_id = start_resp.json()["job_id"]

        # Give it a moment to start
        time.sleep(0.2)

        status_resp = client.get(f"/ingestion/{job_id}/status")
        assert status_resp.status_code == 200
        data = status_resp.json()
        assert data["job_id"] == job_id
        assert data["status"] in ("pending", "running", "completed")

    def test_status_unknown_job(self, client):
        resp = client.get("/ingestion/nonexistent123/status")
        assert resp.status_code == 404


# ---------------------------------------------------------------------------
# POST /ingestion/{job_id}/cancel
# ---------------------------------------------------------------------------

class TestCancelJob:
    def test_cancel_running_job(self, client, input_dir_with_files):
        # Create many files so the job takes a moment
        for i in range(50):
            (input_dir_with_files / f"file_{i:03d}.txt").write_text(f"Content {i}")

        start_resp = client.post("/ingestion/start", json={
            "input_dir": str(input_dir_with_files),
            "output_dir": str(input_dir_with_files.parent / "output"),
            "workers": 1,
        })
        job_id = start_resp.json()["job_id"]

        cancel_resp = client.post(f"/ingestion/{job_id}/cancel")
        assert cancel_resp.status_code == 200
        assert cancel_resp.json()["status"] == "cancelling"

    def test_cancel_unknown_job(self, client):
        resp = client.post("/ingestion/nonexistent123/cancel")
        assert resp.status_code == 404


# ---------------------------------------------------------------------------
# GET /ingestion/{job_id}/stream (SSE)
# ---------------------------------------------------------------------------

class TestSSEStream:
    def test_stream_receives_events(self, client, input_dir_with_files):
        start_resp = client.post("/ingestion/start", json={
            "input_dir": str(input_dir_with_files),
            "output_dir": str(input_dir_with_files.parent / "output"),
            "workers": 1,
        })
        job_id = start_resp.json()["job_id"]

        # Read the SSE stream
        with client.stream("GET", f"/ingestion/{job_id}/stream") as resp:
            assert resp.status_code == 200
            assert "text/event-stream" in resp.headers.get("content-type", "")

            events = []
            for line in resp.iter_lines():
                if line.startswith("event:"):
                    events.append(line)
                if "ingestion_completed" in line:
                    break

            assert len(events) >= 1

    def test_stream_unknown_job(self, client):
        resp = client.get("/ingestion/nonexistent123/stream")
        assert resp.status_code == 404