Files
paperless-ngx/src/documents/tests/test_api_ocr_templates.py
T
shamoon 4ff2cb9c72 Make them owned, duh
Update 0028_ocrtemplate_ocrtemplatezone.py
2026-10-20 11:15:30 -07:00

667 lines
23 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_sets_owner(self):
"""
GIVEN:
- A user creating an OCR template
WHEN:
- API request to create it without an owner
THEN:
- The requesting user owns the template
"""
response = self.client.post(
self.ENDPOINT,
json.dumps(self._make_template_data()),
content_type="application/json",
)
self.assertEqual(response.status_code, status.HTTP_201_CREATED)
self.assertEqual(OcrTemplate.objects.get().owner, self.user)
def test_list_only_shows_permitted_templates(self):
"""
GIVEN:
- Templates that are unowned, owned by the user, owned by another
user, and owned by another user but shared with the user
WHEN:
- The user lists OCR templates
THEN:
- Only the unowned, owned and shared templates are returned
"""
other = User.objects.create_user(username="other")
user = User.objects.create_user(username="template_viewer")
user.user_permissions.add(Permission.objects.get(codename="view_ocrtemplate"))
def template(name, owner=None):
return OcrTemplate.objects.create(
name=name,
document_type=self.doc_type,
source_width=1000,
source_height=1000,
owner=owner,
)
template("Unowned")
template("Mine", owner=user)
template("Private", owner=other)
grant_object(user, template("Shared", owner=other), "view_ocrtemplate")
self.client.force_authenticate(user=user)
response = self.client.get(self.ENDPOINT)
self.assertEqual(response.status_code, status.HTTP_200_OK)
self.assertCountEqual(
[t["name"] for t in response.json()["results"]],
["Unowned", "Mine", "Shared"],
)
def test_cannot_change_another_users_template(self):
"""
GIVEN:
- A template owned by another user
WHEN:
- A user with global change permission tries to update it
THEN:
- The update is refused until the user is granted change on it
"""
other = User.objects.create_user(username="other")
user = User.objects.create_user(username="template_editor")
user.user_permissions.add(
*Permission.objects.filter(
codename__in=["view_ocrtemplate", "change_ocrtemplate"],
),
)
template = OcrTemplate.objects.create(
name="Private",
document_type=self.doc_type,
source_width=1000,
source_height=1000,
owner=other,
)
grant_object(user, template, "view_ocrtemplate")
self.client.force_authenticate(user=user)
url = f"{self.ENDPOINT}{template.pk}/"
response = self.client.patch(url, {"name": "Renamed"}, format="json")
self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
grant_object(user, template, "change_ocrtemplate")
response = self.client.patch(url, {"name": "Renamed"}, format="json")
self.assertEqual(response.status_code, status.HTTP_200_OK)
template.refresh_from_db()
self.assertEqual(template.name, "Renamed")
self.assertEqual(template.owner, other)
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()
@patch("documents.zone_ocr.run_zone_extraction")
def test_run_zone_ocr_requires_global_change_permission(self, extract):
user = User.objects.create_user(username="uploader")
user.user_permissions.add(
*Permission.objects.filter(codename__in=["view_document", "add_document"]),
)
document = DocumentFactory(document_type=self.doc_type, owner=user)
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()