mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2026-10-21 07:18:13 +00:00
561 lines
19 KiB
Python
561 lines
19 KiB
Python
"""Tests for the OCR Template API."""
|
|
|
|
import json
|
|
from unittest.mock import patch
|
|
|
|
from django.contrib.auth.models import Permission
|
|
from django.contrib.auth.models import User
|
|
from rest_framework import status
|
|
from rest_framework.test import APITestCase
|
|
|
|
from documents.models import CustomField
|
|
from documents.models import DocumentType
|
|
from documents.models import OcrTemplate
|
|
from documents.models import OcrTemplateZone
|
|
from paperless_testing.dirs import DirectoriesMixin
|
|
from paperless_testing.factories import DocumentFactory
|
|
from paperless_testing.permissions import grant_object
|
|
|
|
|
|
class TestOcrTemplatesAPI(DirectoriesMixin, APITestCase):
|
|
ENDPOINT = "/api/ocr_templates/"
|
|
|
|
def setUp(self) -> None:
|
|
self.user = User.objects.create_superuser(username="temp_admin")
|
|
self.client.force_authenticate(user=self.user)
|
|
|
|
self.doc_type = DocumentType.objects.create(name="Invoice")
|
|
self.custom_field_text = CustomField.objects.create(
|
|
name="Invoice Number",
|
|
data_type=CustomField.FieldDataType.STRING,
|
|
)
|
|
self.custom_field_date = CustomField.objects.create(
|
|
name="Invoice Date",
|
|
data_type=CustomField.FieldDataType.DATE,
|
|
)
|
|
self.custom_field_int = CustomField.objects.create(
|
|
name="Amount",
|
|
data_type=CustomField.FieldDataType.INT,
|
|
)
|
|
self.custom_field_doclink = CustomField.objects.create(
|
|
name="Related Docs",
|
|
data_type=CustomField.FieldDataType.DOCUMENTLINK,
|
|
)
|
|
|
|
return super().setUp()
|
|
|
|
def _make_template_data(self, **overrides):
|
|
data = {
|
|
"name": "Invoice Template",
|
|
"document_type": self.doc_type.pk,
|
|
"default_page": 0,
|
|
"source_width": 2480,
|
|
"source_height": 3508,
|
|
"enabled": True,
|
|
"zones": [],
|
|
}
|
|
data.update(overrides)
|
|
return data
|
|
|
|
def _make_zone_data(self, **overrides):
|
|
data = {
|
|
"name": "Zone 1",
|
|
"custom_field": self.custom_field_text.pk,
|
|
"x": 100,
|
|
"y": 100,
|
|
"width": 200,
|
|
"height": 50,
|
|
"ocr_language": "deu+eng",
|
|
"transform": "strip",
|
|
"order": 0,
|
|
}
|
|
data.update(overrides)
|
|
return data
|
|
|
|
# --- Create ---
|
|
|
|
def test_create_template(self):
|
|
"""
|
|
GIVEN:
|
|
- A document type and custom fields exist
|
|
WHEN:
|
|
- API request to create an OCR template with one zone
|
|
THEN:
|
|
- The template and zone are created
|
|
"""
|
|
data = self._make_template_data(
|
|
zones=[
|
|
self._make_zone_data(
|
|
name="Invoice Number",
|
|
x=1500,
|
|
y=200,
|
|
width=800,
|
|
height=100,
|
|
),
|
|
],
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
|
|
|
|
result = resp.json()
|
|
self.assertEqual(result["name"], "Invoice Template")
|
|
self.assertEqual(result["document_type"], self.doc_type.pk)
|
|
self.assertEqual(len(result["zones"]), 1)
|
|
self.assertEqual(result["zones"][0]["name"], "Invoice Number")
|
|
self.assertEqual(OcrTemplate.objects.count(), 1)
|
|
self.assertEqual(OcrTemplateZone.objects.count(), 1)
|
|
|
|
def test_create_template_multiple_zones(self):
|
|
"""
|
|
GIVEN:
|
|
- Multiple custom fields exist
|
|
WHEN:
|
|
- A template with multiple zones is created
|
|
THEN:
|
|
- All zones are created
|
|
"""
|
|
data = self._make_template_data(
|
|
zones=[
|
|
self._make_zone_data(
|
|
name="Invoice Number",
|
|
custom_field=self.custom_field_text.pk,
|
|
),
|
|
self._make_zone_data(
|
|
name="Invoice Date",
|
|
custom_field=self.custom_field_date.pk,
|
|
order=1,
|
|
),
|
|
],
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
|
|
self.assertEqual(len(resp.json()["zones"]), 2)
|
|
self.assertEqual(OcrTemplateZone.objects.count(), 2)
|
|
|
|
def test_create_template_no_zones(self):
|
|
"""
|
|
GIVEN:
|
|
- Valid template data without zones
|
|
WHEN:
|
|
- Template is created
|
|
THEN:
|
|
- Template is created with no zones
|
|
"""
|
|
data = self._make_template_data()
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_201_CREATED)
|
|
self.assertEqual(len(resp.json()["zones"]), 0)
|
|
|
|
# --- Validation ---
|
|
|
|
def test_create_template_zero_source_width_rejected(self):
|
|
"""
|
|
GIVEN:
|
|
- Template data with source_width=0
|
|
WHEN:
|
|
- Create is attempted
|
|
THEN:
|
|
- 400 error is returned
|
|
"""
|
|
data = self._make_template_data(source_width=0)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_create_template_zero_source_height_rejected(self):
|
|
data = self._make_template_data(source_height=0)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_create_zone_zero_width_rejected(self):
|
|
data = self._make_template_data(
|
|
zones=[self._make_zone_data(width=0)],
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_create_zone_zero_height_rejected(self):
|
|
data = self._make_template_data(
|
|
zones=[self._make_zone_data(height=0)],
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_create_zone_exceeds_source_width_rejected(self):
|
|
"""Zone that extends beyond the source image width should be rejected."""
|
|
data = self._make_template_data(
|
|
source_width=1000,
|
|
zones=[self._make_zone_data(x=800, width=300)], # 800+300 > 1000
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_create_zone_exceeds_source_height_rejected(self):
|
|
data = self._make_template_data(
|
|
source_height=1000,
|
|
zones=[self._make_zone_data(y=900, height=200)], # 900+200 > 1000
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
def test_create_zone_unsupported_custom_field_type_rejected(self):
|
|
"""DOCUMENTLINK and SELECT fields can't be populated via OCR."""
|
|
data = self._make_template_data(
|
|
zones=[self._make_zone_data(custom_field=self.custom_field_doclink.pk)],
|
|
)
|
|
resp = self.client.post(
|
|
self.ENDPOINT,
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_400_BAD_REQUEST)
|
|
|
|
# --- List ---
|
|
|
|
def test_list_templates(self):
|
|
template = OcrTemplate.objects.create(
|
|
name="Test Template",
|
|
document_type=self.doc_type,
|
|
source_width=2480,
|
|
source_height=3508,
|
|
)
|
|
OcrTemplateZone.objects.create(
|
|
template=template,
|
|
name="Zone 1",
|
|
custom_field=self.custom_field_text,
|
|
x=100,
|
|
y=100,
|
|
width=200,
|
|
height=50,
|
|
)
|
|
|
|
resp = self.client.get(self.ENDPOINT)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
data = resp.json()
|
|
self.assertEqual(data["count"], 1)
|
|
self.assertEqual(len(data["results"][0]["zones"]), 1)
|
|
|
|
def test_list_empty(self):
|
|
resp = self.client.get(self.ENDPOINT)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(resp.json()["count"], 0)
|
|
|
|
# --- Update ---
|
|
|
|
def test_update_template_replaces_zones(self):
|
|
"""PUT should replace all zones with the new set."""
|
|
template = OcrTemplate.objects.create(
|
|
name="Old Name",
|
|
document_type=self.doc_type,
|
|
source_width=2480,
|
|
source_height=3508,
|
|
)
|
|
OcrTemplateZone.objects.create(
|
|
template=template,
|
|
name="Old Zone",
|
|
custom_field=self.custom_field_text,
|
|
x=0,
|
|
y=0,
|
|
width=100,
|
|
height=100,
|
|
)
|
|
|
|
data = self._make_template_data(
|
|
name="New Name",
|
|
zones=[
|
|
self._make_zone_data(
|
|
name="New Zone",
|
|
custom_field=self.custom_field_date.pk,
|
|
),
|
|
],
|
|
)
|
|
resp = self.client.put(
|
|
f"{self.ENDPOINT}{template.pk}/",
|
|
data=json.dumps(data),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
|
|
template.refresh_from_db()
|
|
self.assertEqual(template.name, "New Name")
|
|
self.assertEqual(OcrTemplateZone.objects.count(), 1)
|
|
self.assertEqual(OcrTemplateZone.objects.first().name, "New Zone")
|
|
|
|
# --- Delete ---
|
|
|
|
def test_delete_template_cascades_zones(self):
|
|
template = OcrTemplate.objects.create(
|
|
name="To Delete",
|
|
document_type=self.doc_type,
|
|
source_width=2480,
|
|
source_height=3508,
|
|
)
|
|
OcrTemplateZone.objects.create(
|
|
template=template,
|
|
name="Zone",
|
|
custom_field=self.custom_field_text,
|
|
x=0,
|
|
y=0,
|
|
width=100,
|
|
height=100,
|
|
)
|
|
|
|
resp = self.client.delete(f"{self.ENDPOINT}{template.pk}/")
|
|
self.assertEqual(resp.status_code, status.HTTP_204_NO_CONTENT)
|
|
self.assertEqual(OcrTemplate.objects.count(), 0)
|
|
self.assertEqual(OcrTemplateZone.objects.count(), 0)
|
|
|
|
def test_delete_nonexistent_returns_404(self):
|
|
resp = self.client.delete(f"{self.ENDPOINT}99999/")
|
|
self.assertEqual(resp.status_code, status.HTTP_404_NOT_FOUND)
|
|
|
|
# --- Patch ---
|
|
|
|
def test_patch_toggle_enabled(self):
|
|
template = OcrTemplate.objects.create(
|
|
name="Toggle Test",
|
|
document_type=self.doc_type,
|
|
source_width=2480,
|
|
source_height=3508,
|
|
enabled=True,
|
|
)
|
|
|
|
resp = self.client.patch(
|
|
f"{self.ENDPOINT}{template.pk}/",
|
|
data=json.dumps({"enabled": False}),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
template.refresh_from_db()
|
|
self.assertFalse(template.enabled)
|
|
|
|
def test_patch_preserves_zones(self):
|
|
"""PATCH without zones field should not delete existing zones."""
|
|
template = OcrTemplate.objects.create(
|
|
name="Patch Test",
|
|
document_type=self.doc_type,
|
|
source_width=2480,
|
|
source_height=3508,
|
|
)
|
|
OcrTemplateZone.objects.create(
|
|
template=template,
|
|
name="Existing Zone",
|
|
custom_field=self.custom_field_text,
|
|
x=0,
|
|
y=0,
|
|
width=100,
|
|
height=100,
|
|
)
|
|
|
|
resp = self.client.patch(
|
|
f"{self.ENDPOINT}{template.pk}/",
|
|
data=json.dumps({"name": "Updated Name"}),
|
|
content_type="application/json",
|
|
)
|
|
self.assertEqual(resp.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(OcrTemplateZone.objects.count(), 1)
|
|
|
|
# --- Auth ---
|
|
|
|
def test_unauthenticated_rejected(self):
|
|
self.client.logout()
|
|
resp = self.client.get(self.ENDPOINT)
|
|
self.assertIn(
|
|
resp.status_code,
|
|
(status.HTTP_401_UNAUTHORIZED, status.HTTP_403_FORBIDDEN),
|
|
)
|
|
|
|
@patch("documents.zone_ocr.run_zone_extraction")
|
|
def test_run_zone_ocr_without_document_type_is_skipped(self, extract):
|
|
document = DocumentFactory()
|
|
|
|
response = self.client.post(f"/api/documents/{document.pk}/run-zone-ocr/")
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(
|
|
response.json(),
|
|
{"status": "skipped", "reason": "no_document_type", "results": []},
|
|
)
|
|
extract.assert_not_called()
|
|
|
|
@patch("documents.zone_ocr.run_zone_extraction")
|
|
def test_run_zone_ocr_without_enabled_template_is_skipped(self, extract):
|
|
document = DocumentFactory(document_type=self.doc_type)
|
|
|
|
for enabled in (None, False):
|
|
with self.subTest(enabled=enabled):
|
|
if enabled is False:
|
|
OcrTemplate.objects.create(
|
|
name="Disabled template",
|
|
document_type=self.doc_type,
|
|
source_width=1000,
|
|
source_height=1000,
|
|
enabled=False,
|
|
)
|
|
|
|
response = self.client.post(
|
|
f"/api/documents/{document.pk}/run-zone-ocr/",
|
|
)
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(
|
|
response.json(),
|
|
{"status": "skipped", "reason": "no_template", "results": []},
|
|
)
|
|
extract.assert_not_called()
|
|
|
|
def test_document_page_image_requires_view_permission_on_the_document(self):
|
|
owner = User.objects.create_user(username="owner")
|
|
user = User.objects.create_user(username="template_editor")
|
|
user.user_permissions.add(Permission.objects.get(codename="view_ocrtemplate"))
|
|
document = DocumentFactory(owner=owner)
|
|
self.client.force_authenticate(user=user)
|
|
url = f"{self.ENDPOINT}document-page-image/{document.pk}/0/"
|
|
|
|
response = self.client.get(url)
|
|
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
|
|
|
grant_object(user, document, "view_document")
|
|
response = self.client.get(url)
|
|
# Past the permission check; the factory document has no file on disk.
|
|
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
|
|
|
@patch("documents.zone_ocr.extract_zone_preview")
|
|
def test_test_zone_requires_view_permission_on_the_document(self, preview):
|
|
owner = User.objects.create_user(username="owner")
|
|
user = User.objects.create_user(username="template_editor")
|
|
user.user_permissions.add(Permission.objects.get(codename="add_ocrtemplate"))
|
|
document = DocumentFactory(owner=owner)
|
|
self.client.force_authenticate(user=user)
|
|
data = {
|
|
"document": document.pk,
|
|
"zone": {"x": 0, "y": 0, "width": 100, "height": 100, "page": 1},
|
|
}
|
|
|
|
response = self.client.post(f"{self.ENDPOINT}test-zone/", data, format="json")
|
|
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
|
|
|
grant_object(user, document, "view_document")
|
|
response = self.client.post(f"{self.ENDPOINT}test-zone/", data, format="json")
|
|
# Past the permission check; the factory document has no file on disk.
|
|
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
|
preview.assert_not_called()
|
|
|
|
@patch("documents.zone_ocr.run_zone_extraction")
|
|
def test_run_zone_ocr_requires_change_permission_on_the_document(
|
|
self,
|
|
extract,
|
|
):
|
|
owner = User.objects.create_user(username="owner")
|
|
user = User.objects.create_user(username="editor")
|
|
user.user_permissions.add(
|
|
*Permission.objects.filter(
|
|
codename__in=["view_document", "change_document"],
|
|
),
|
|
)
|
|
document = DocumentFactory(document_type=self.doc_type, owner=owner)
|
|
self.client.force_authenticate(user=user)
|
|
|
|
response = self.client.post(f"/api/documents/{document.pk}/run-zone-ocr/")
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
|
|
extract.assert_not_called()
|
|
|
|
def test_run_zone_ocr_missing_document_is_an_error(self):
|
|
response = self.client.post("/api/documents/99999/run-zone-ocr/")
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
|
|
|
@patch("documents.zone_ocr.run_zone_extraction")
|
|
def test_run_zone_ocr_missing_file_is_an_error(self, extract):
|
|
document = DocumentFactory(document_type=self.doc_type)
|
|
OcrTemplate.objects.create(
|
|
name="Template",
|
|
document_type=self.doc_type,
|
|
source_width=1000,
|
|
source_height=1000,
|
|
)
|
|
|
|
response = self.client.post(f"/api/documents/{document.pk}/run-zone-ocr/")
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
|
|
self.assertEqual(response.json(), {"error": "Document file not found"})
|
|
extract.assert_not_called()
|
|
|
|
@patch("documents.zone_ocr.run_zone_extraction")
|
|
def test_run_zone_ocr_reports_completed_results(self, extract):
|
|
document = DocumentFactory(
|
|
document_type=self.doc_type,
|
|
mime_type="application/pdf",
|
|
)
|
|
document.source_path.write_bytes(b"test PDF")
|
|
template = OcrTemplate.objects.create(
|
|
name="Template",
|
|
document_type=self.doc_type,
|
|
source_width=1000,
|
|
source_height=1000,
|
|
)
|
|
OcrTemplateZone.objects.create(
|
|
template=template,
|
|
name="Title zone",
|
|
target="title",
|
|
x=0,
|
|
y=0,
|
|
width=100,
|
|
height=100,
|
|
)
|
|
|
|
response = self.client.post(f"/api/documents/{document.pk}/run-zone-ocr/")
|
|
|
|
self.assertEqual(response.status_code, status.HTTP_200_OK)
|
|
self.assertEqual(
|
|
response.json(),
|
|
{
|
|
"status": "completed",
|
|
"results": [
|
|
{
|
|
"template": template.name,
|
|
"zone": "Title zone",
|
|
"custom_field": "Title",
|
|
"value": document.title,
|
|
},
|
|
],
|
|
},
|
|
)
|
|
extract.assert_called_once()
|