Files
paperless-ngx/src/documents/tests/utils.py
T
stumpylog 8849aff5d9 Chore: Run read-only migration tests through one migration cycle per class
Each TestMigrations test migrates back, seeds, migrates forward and returns to
the latest migration, which is slow. The fulltext query
prefix, sha256 checksum and skip archive file classes only read the migrated
data, so they now opt in to migrate_once: the first test runs the migration,
the rest reuse its state, and the schema is restored and the tables flushed in
tearDownClass.
2026-09-21 11:45:28 -07:00

390 lines
13 KiB
Python

import time
import warnings
from collections.abc import Callable
from collections.abc import Generator
from collections.abc import Iterator
from contextlib import contextmanager
from os import PathLike
from pathlib import Path
from typing import Any
from unittest import mock
import httpx
import pytest
from django.apps import apps
from django.db import connection
from django.db.migrations.executor import MigrationExecutor
from django.http import StreamingHttpResponse
from django.test import TransactionTestCase
from documents.consumer import AsnCheckPlugin
from documents.consumer import ConsumerPlugin
from documents.consumer import ConsumerPreflightPlugin
from documents.data_models import ConsumableDocument
from documents.data_models import DocumentMetadataOverrides
from documents.data_models import DocumentSource
from documents.parsers import ParseError
from documents.plugins.helpers import ProgressStatusOptions
def util_call_with_backoff(
method_or_callable: Callable,
args: list | tuple,
*,
skip_on_50x_err=True,
) -> tuple[bool, Any]:
"""
For whatever reason, the images started during the test pipeline like to
segfault sometimes, crash and otherwise fail randomly, when run with the
exact files that usually pass.
So, this function will retry the given method/function up to 3 times, with larger backoff
periods between each attempt, in hopes the issue resolves itself during
one attempt to parse.
This will wait the following:
- Attempt 1 - 20s following failure
- Attempt 2 - 40s following failure
- Attempt 3 - 80s following failure
"""
result = None
succeeded = False
retry_time = 20.0
retry_count = 0
status_codes = []
max_retry_count = 3
while retry_count < max_retry_count and not succeeded:
try:
result = method_or_callable(*args)
succeeded = True
except ParseError as e: # pragma: no cover
cause_exec = e.__cause__
if cause_exec is not None and isinstance(cause_exec, httpx.HTTPStatusError):
status_codes.append(cause_exec.response.status_code)
warnings.warn(
f"HTTP Exception for {cause_exec.request.url} - {cause_exec}",
)
else:
warnings.warn(f"Unexpected error: {e}")
except Exception as e: # pragma: no cover
warnings.warn(f"Unexpected error: {e}")
retry_count = retry_count + 1
time.sleep(retry_time)
retry_time = retry_time * 2.0
if (
not succeeded
and status_codes
and skip_on_50x_err
and all(httpx.codes.is_server_error(code) for code in status_codes)
):
pytest.skip("Repeated HTTP 50x for service") # pragma: no cover
return succeeded, result
def read_streaming_response(response: StreamingHttpResponse) -> bytes:
"""Consume a StreamingHttpResponse/FileResponse and close it."""
content = b"".join(response.streaming_content)
response.close()
return content
class FileSystemAssertsMixin:
"""
Utilities for checks various state information of the file system
"""
def assertIsFile(self, path: PathLike[str] | str) -> None:
self.assertTrue(Path(path).resolve().is_file(), f"File does not exist: {path}")
def assertIsNotFile(self, path: PathLike[str] | str) -> None:
self.assertFalse(Path(path).resolve().is_file(), f"File does exist: {path}")
def assertIsDir(self, path: PathLike[str] | str) -> None:
self.assertTrue(Path(path).resolve().is_dir(), f"Dir does not exist: {path}")
def assertIsNotDir(self, path: PathLike[str] | str) -> None:
self.assertFalse(Path(path).resolve().is_dir(), f"Dir does exist: {path}")
def assertFileCountInDir(self, path: PathLike[str] | str, count: int) -> None:
path = Path(path).resolve()
self.assertTrue(path.is_dir(), f"Path {path} is not a directory")
files = [x for x in path.iterdir() if x.is_file()]
self.assertEqual(
len(files),
count,
f"Path {path} contains {len(files)} files instead of {count} files",
)
class ConsumeTaskMixin:
"""
Provides mocking of the consume_file asynchronous task and useful utilities
for decoding its arguments
"""
def setUp(self) -> None:
self.consume_file_patcher = mock.patch(
"documents.tasks.consume_file.apply_async",
)
self.consume_file_mock = self.consume_file_patcher.start()
super().setUp()
def tearDown(self) -> None:
super().tearDown()
self.consume_file_patcher.stop()
def assert_queue_consumption_task_call_args(
self,
) -> tuple[ConsumableDocument, DocumentMetadataOverrides]:
"""Assert the task was queued exactly once and return its call args."""
self.consume_file_mock.assert_called_once()
task_kwargs = self.consume_file_mock.call_args.kwargs["kwargs"]
return (task_kwargs["input_doc"], task_kwargs["overrides"])
def get_all_consume_task_call_args(
self,
) -> Iterator[tuple[ConsumableDocument, DocumentMetadataOverrides]]:
"""Iterate over all queued consume task calls and yield their call args."""
self.consume_file_mock.assert_called()
for call in self.consume_file_mock.call_args_list:
task_kwargs = call.kwargs["kwargs"]
yield (task_kwargs["input_doc"], task_kwargs["overrides"])
class TestMigrations(TransactionTestCase):
"""Run a migration on seeded data, then let the tests inspect the result.
By default every test migrates back, seeds, migrates forward and returns to
the latest migration, which costs several seconds. A class whose tests only
read the migrated data can set ``migrate_once`` to pay that once per class:
the migration runs for the first test, the database is left alone between
tests, and it is restored and flushed when the class finishes. Such tests
must not write to the database.
"""
@property
def app(self):
return apps.get_containing_app_config(type(self).__module__).name
migrate_from = None
dependencies = None
migrate_to = None
migrate_once = False
_once_owner: "TestMigrations | None" = None
_once_state: dict[str, Any] | None = None
_once_finishing = False
def setUp(self) -> None:
super().setUp()
cls = type(self)
if self.migrate_once:
if cls._once_state is not None:
vars(self).update(cls._once_state)
return
if cls._once_owner is not None:
raise RuntimeError(
f"The migration in '{cls.__name__}' failed for an earlier test",
)
# Recorded before migrating so a failed migration is still restored
cls._once_owner = self
before = dict(vars(self))
self._migrate()
if self.migrate_once:
cls._once_state = {
name: value
for name, value in vars(self).items()
if name not in before or before[name] is not value
}
def _migrate(self) -> None:
assert self.migrate_from and self.migrate_to, (
f"TestCase '{type(self).__name__}' must define migrate_from and migrate_to properties"
)
self.migrate_from = [(self.app, self.migrate_from)]
if self.dependencies is not None:
self.migrate_from.extend(self.dependencies)
self.migrate_to = [(self.app, self.migrate_to)]
executor = MigrationExecutor(connection)
old_apps = executor.loader.project_state(self.migrate_from).apps
# Reverse to the original migration
executor.migrate(self.migrate_from)
self.setUpBeforeMigration(old_apps)
self.apps = old_apps
# Run the migration to test
executor = MigrationExecutor(connection)
executor.loader.build_graph() # reload.
executor.migrate(self.migrate_to)
self.apps = executor.loader.project_state(self.migrate_to).apps
def setUpBeforeMigration(self, apps) -> None:
pass
def _migrate_to_latest(self) -> None:
executor = MigrationExecutor(connection)
executor.loader.build_graph()
targets = executor.loader.graph.leaf_nodes()
executor.migrate(targets)
def tearDown(self) -> None:
"""
Ensure the database schema is restored to the latest migration after
each migration test, so subsequent tests run against HEAD.
"""
if self.migrate_once and not self._once_finishing:
return
try:
self._migrate_to_latest()
finally:
super().tearDown()
def _fixture_teardown(self) -> None:
# Django flushes every table after each test, which would discard the
# data the remaining tests of a migrate_once class still need
if self.migrate_once and not self._once_finishing:
return
super()._fixture_teardown()
@classmethod
def tearDownClass(cls) -> None:
owner = cls._once_owner
try:
if owner is not None:
cls._once_finishing = True
owner.tearDown()
owner._fixture_teardown()
finally:
cls._once_owner = None
cls._once_state = None
cls._once_finishing = False
super().tearDownClass()
class SampleDirMixin:
SAMPLE_DIR = Path(__file__).parent / "samples"
BARCODE_SAMPLE_DIR = SAMPLE_DIR / "barcodes"
class GetConsumerMixin:
@contextmanager
def get_consumer(
self,
filepath: Path,
overrides: DocumentMetadataOverrides | None = None,
source: DocumentSource = DocumentSource.ConsumeFolder,
mailrule_id: int | None = None,
) -> Generator[ConsumerPlugin, None, None]:
# Store this for verification
self.status = DummyProgressManager(filepath.name, None)
doc = ConsumableDocument(
source,
original_file=filepath,
mailrule_id=mailrule_id or None,
)
preflight_plugin = ConsumerPreflightPlugin(
doc,
overrides or DocumentMetadataOverrides(),
self.status, # type: ignore
self.dirs.scratch_dir,
"task-id",
)
preflight_plugin.setup()
asncheck_plugin = AsnCheckPlugin(
doc,
overrides or DocumentMetadataOverrides(),
self.status, # type: ignore
self.dirs.scratch_dir,
"task-id",
)
asncheck_plugin.setup()
reader = ConsumerPlugin(
doc,
overrides or DocumentMetadataOverrides(),
self.status, # type: ignore
self.dirs.scratch_dir,
"task-id",
)
reader.setup()
try:
preflight_plugin.run()
asncheck_plugin.run()
yield reader
finally:
reader.cleanup()
class DummyProgressManager:
"""
A dummy handler for progress management that doesn't actually try to
connect to Redis. Payloads are stored for test assertions if needed.
Use it with
mock.patch("documents.tasks.ProgressManager", DummyProgressManager)
"""
def __init__(self, filename: str, task_id: str | None = None) -> None:
self.filename = filename
self.task_id = task_id
self.payloads = []
def __enter__(self):
self.open()
return self
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
self.close()
def open(self) -> None:
pass
def close(self) -> None:
pass
def send_progress(
self,
status: ProgressStatusOptions,
message: str,
current_progress: int,
max_progress: int,
*,
document_id: int | None = None,
owner_id: int | None = None,
users_can_view: list[int] | None = None,
groups_can_view: list[int] | None = None,
) -> None:
# Ensure the layer is open
self.open()
payload = {
"type": "status_update",
"data": {
"filename": self.filename,
"task_id": self.task_id,
"current_progress": current_progress,
"max_progress": max_progress,
"status": status,
"message": message,
"document_id": document_id,
"owner_id": owner_id,
"users_can_view": users_can_view or [],
"groups_can_view": groups_can_view or [],
},
}
self.payloads.append(payload)