"""Tests for cf-docuvision — mock inference path only (no GPU required)."""
import pytest
from unittest.mock import MagicMock, patch
from app.dolphin import (
DolphinElement,
DolphinParser,
DolphinResult,
dolphin_to_cf_elements,
_TYPE_MAP,
)
class TestTypeMap:
def test_all_21_dolphin_types_mapped(self):
# Spot-check the 21 Dolphin-v2 element types are all covered
expected = {
"title", "section_header", "paragraph", "caption", "footnote",
"page_header", "page_footer", "list_item", "table", "figure",
"figure_caption", "formula", "code", "annotation", "abstract",
"toc_item", "reference", "equation", "watermark", "stamp", "signature",
}
assert expected == set(_TYPE_MAP.keys())
def test_table_maps_to_table(self):
assert _TYPE_MAP["table"] == "table"
def test_title_maps_to_heading(self):
assert _TYPE_MAP["title"] == "heading"
def test_formula_maps_to_formula(self):
assert _TYPE_MAP["formula"] == "formula"
class TestDolphinToCfElements:
def _make_result(self, elements: list[DolphinElement]) -> DolphinResult:
raw_text = "\n".join(e.text for e in elements if e.text)
return DolphinResult(elements=elements, raw_text=raw_text)
def test_paragraph_goes_to_elements(self):
result = self._make_result([DolphinElement("paragraph", "Hello world")])
elements, tables = dolphin_to_cf_elements(result)
assert len(elements) == 1
assert elements[0]["type"] == "paragraph"
assert elements[0]["text"] == "Hello world"
assert tables == []
def test_table_with_html_goes_to_tables(self):
result = self._make_result([
DolphinElement("table", "col1 col2", html="
")
])
elements, tables = dolphin_to_cf_elements(result)
assert len(tables) == 1
assert "" in tables[0]["html"]
assert elements == []
def test_table_without_html_goes_to_elements(self):
result = self._make_result([DolphinElement("table", "some table text", html=None)])
elements, tables = dolphin_to_cf_elements(result)
assert len(elements) == 1
assert tables == []
def test_bbox_preserved(self):
result = self._make_result([
DolphinElement("paragraph", "text", bbox=[0.1, 0.2, 0.8, 0.3])
])
elements, _ = dolphin_to_cf_elements(result)
assert elements[0]["bbox"] == [0.1, 0.2, 0.8, 0.3]
def test_mixed_elements_and_tables(self):
result = self._make_result([
DolphinElement("title", "Document Title"),
DolphinElement("table", "data", html=""),
DolphinElement("paragraph", "Body text"),
])
elements, tables = dolphin_to_cf_elements(result)
assert len(elements) == 2
assert len(tables) == 1
assert elements[0]["type"] == "heading"
def test_empty_result(self):
result = DolphinResult()
elements, tables = dolphin_to_cf_elements(result)
assert elements == []
assert tables == []
class TestParseOutputFallbacks:
"""Test _parse_dolphin_output without loading the real model."""
def _make_parser(self) -> DolphinParser:
"""Create a DolphinParser without loading the model."""
parser = object.__new__(DolphinParser)
parser._model_id = "ByteDance/Dolphin-v2"
parser._device = "cpu"
return parser
def test_json_array_output(self):
parser = self._make_parser()
raw = '[{"type": "paragraph", "text": "Hello", "bbox": null, "html": null}]'
elements = parser._parse_dolphin_output(raw)
assert len(elements) == 1
assert elements[0]["type"] == "paragraph"
def test_element_tag_output(self):
parser = self._make_parser()
raw = 'My Title'
elements = parser._parse_dolphin_output(raw)
assert len(elements) == 1
assert elements[0]["type"] == "title"
assert elements[0]["text"] == "My Title"
assert elements[0]["bbox"] == [0.1, 0.2, 0.8, 0.3]
def test_element_tag_without_bbox(self):
parser = self._make_parser()
raw = 'Plain text'
elements = parser._parse_dolphin_output(raw)
assert elements[0]["bbox"] is None
def test_fallback_to_single_paragraph(self):
parser = self._make_parser()
raw = "This is some unstructured text output."
elements = parser._parse_dolphin_output(raw)
assert len(elements) == 1
assert elements[0]["type"] == "paragraph"
assert "unstructured text" in elements[0]["text"]
def test_empty_output(self):
parser = self._make_parser()
elements = parser._parse_dolphin_output("")
assert elements == []
class TestFastAPIRoutes:
"""Integration tests for the FastAPI endpoints using TestClient."""
def _make_app_with_mock_parser(self):
from fastapi.testclient import TestClient
import app.main as main_module
mock_parser = MagicMock()
mock_parser._model_id = "ByteDance/Dolphin-v2"
from app.dolphin import DolphinResult, DolphinElement
mock_result = DolphinResult(
elements=[DolphinElement("paragraph", "Extracted text")],
raw_text="Extracted text",
)
mock_parser.parse_b64.return_value = mock_result
main_module._parser = mock_parser
return TestClient(main_module.app)
def test_health_with_loaded_model(self):
client = self._make_app_with_mock_parser()
resp = client.get("/health")
assert resp.status_code == 200
assert resp.json()["status"] == "ok"
def test_health_without_model_returns_503(self):
from fastapi.testclient import TestClient
import app.main as main_module
main_module._parser = None
client = TestClient(main_module.app, raise_server_exceptions=False)
resp = client.get("/health")
assert resp.status_code == 503
def test_extract_returns_structured_response(self):
import base64
client = self._make_app_with_mock_parser()
payload = {
"image_b64": base64.b64encode(b"fake-image-bytes").decode(),
"hint": "auto",
}
resp = client.post("/extract", json=payload)
assert resp.status_code == 200
data = resp.json()
assert "elements" in data
assert "tables" in data
assert "raw_text" in data
assert data["metadata"]["source"] == "cf-docuvision"
def test_extract_invalid_hint_returns_422(self):
import base64
client = self._make_app_with_mock_parser()
payload = {
"image_b64": base64.b64encode(b"fake").decode(),
"hint": "invalid",
}
resp = client.post("/extract", json=payload)
assert resp.status_code == 422