Stirling-PDF/engine/tests/test_document_classifier_routes.py

from __future__ import annotations

from collections.abc import Iterator

import pytest
from fastapi.testclient import TestClient

from stirling.api import app
from stirling.api.dependencies import get_document_classifier_agent
from stirling.contracts import (
    ClassifyDocumentRequest,
    ClassifyDocumentResponse,
    DocumentClassificationResponse,
)


class StubClassifierAgent:
    """Stands in for DocumentClassifierAgent so route tests don't call a model."""

    def __init__(self, response: ClassifyDocumentResponse) -> None:
        self._response = response

    async def classify(self, _request: ClassifyDocumentRequest) -> ClassifyDocumentResponse:
        return self._response


@pytest.fixture
def classification_client() -> Iterator[TestClient]:
    app.dependency_overrides[get_document_classifier_agent] = lambda: StubClassifierAgent(
        # The result is label ids (the model's name answers, mapped to ids).
        DocumentClassificationResponse(labels=["nda", "contract"])
    )
    try:
        yield TestClient(app)
    finally:
        app.dependency_overrides.pop(get_document_classifier_agent, None)


_LABELS = [{"id": "nda", "name": "NDA"}, {"id": "contract", "name": "Contract"}]


def test_classify_returns_assigned_labels(classification_client: TestClient) -> None:
    response = classification_client.post(
        "/api/v1/documents/classify",
        json={
            "fileName": "nda.pdf",
            "pages": [{"pageNumber": 1, "text": "Mutual NDA between A and B."}],
            "labels": _LABELS,
        },
    )
    assert response.status_code == 200
    assert response.json() == {"labels": ["nda", "contract"]}


def test_classify_accepts_allowed_labels_on_the_request(classification_client: TestClient) -> None:
    response = classification_client.post(
        "/api/v1/documents/classify",
        json={
            "fileName": "nda.pdf",
            "pages": [],
            "labels": [
                {"id": "contract", "name": "Contract"},
                {"id": "invoice", "name": "Invoice"},
            ],
        },
    )
    assert response.status_code == 200


def test_classify_accepts_empty_pages(classification_client: TestClient) -> None:
    response = classification_client.post(
        "/api/v1/documents/classify",
        json={"fileName": "blank.pdf", "pages": [], "labels": _LABELS},
    )
    assert response.status_code == 200


def test_classify_rejects_empty_file_name(classification_client: TestClient) -> None:
    response = classification_client.post(
        "/api/v1/documents/classify",
        json={"fileName": "", "pages": [], "labels": _LABELS},
    )
    assert response.status_code == 422


def test_classify_rejects_missing_labels(classification_client: TestClient) -> None:
    # The backend always supplies the vocabulary; a request without one is invalid.
    response = classification_client.post(
        "/api/v1/documents/classify",
        json={"fileName": "nda.pdf", "pages": []},
    )
    assert response.status_code == 422