wehub-resource-sync
93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
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
|