import pickle import re import warnings from datetime import UTC from datetime import datetime from pathlib import Path from unittest import mock import numpy as np import pytest from django.conf import settings from django.db import connection from django.test import TestCase from django.test import override_settings from django.test.utils import CaptureQueriesContext from pytest_django.fixtures import Settings from pytest_mock import MockerFixture from documents.classifier import ClassifierModelCorruptError from documents.classifier import DocumentClassifier from documents.classifier import IncompatibleClassifierVersionError from documents.classifier import _predict_with_threshold from documents.classifier import _text_analyzer from documents.classifier import load_classifier from documents.models import Correspondent from documents.models import Document from documents.models import DocumentType from documents.models import MatchingModel from documents.models import StoragePath from documents.models import Tag from paperless.settings import CLASSIFIER_LANGUAGES from paperless.signed_pickle import HMAC_SIZE from paperless.signed_pickle import signed_pickle_dumps from paperless_testing.dirs import DirectoriesMixin from paperless_testing.factories import DocumentFactory from paperless_testing.factories import TagFactory def dummy_preprocess(content: str) -> str: """ Simpler, faster pre-processing for testing purposes """ content = content.lower().strip() content = re.sub(r"\s+", " ", content) return content class TestClassifier(DirectoriesMixin, TestCase): def setUp(self) -> None: super().setUp() self.classifier = DocumentClassifier() self.classifier.preprocess_content = mock.MagicMock( side_effect=dummy_preprocess, ) def generate_test_data(self) -> None: self.c1 = Correspondent.objects.create( name="c1", matching_algorithm=Correspondent.MATCH_AUTO, ) self.c2 = Correspondent.objects.create(name="c2") self.c3 = Correspondent.objects.create( name="c3", matching_algorithm=Correspondent.MATCH_AUTO, ) self.t1 = Tag.objects.create( name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12, ) self.t2 = Tag.objects.create( name="t2", matching_algorithm=Tag.MATCH_ANY, pk=34, is_inbox_tag=True, ) self.t3 = Tag.objects.create( name="t3", matching_algorithm=Tag.MATCH_AUTO, pk=45, ) self.t4 = Tag.objects.create( name="t4", matching_algorithm=Tag.MATCH_ANY, pk=46, ) self.dt = DocumentType.objects.create( name="dt", matching_algorithm=DocumentType.MATCH_AUTO, ) self.dt2 = DocumentType.objects.create( name="dt2", matching_algorithm=DocumentType.MATCH_AUTO, ) self.sp1 = StoragePath.objects.create( name="sp1", path="path1", matching_algorithm=DocumentType.MATCH_AUTO, ) self.sp2 = StoragePath.objects.create( name="sp2", path="path2", matching_algorithm=DocumentType.MATCH_AUTO, ) self.store_paths = [self.sp1, self.sp2] self.doc1 = Document.objects.create( title="doc1", content="this is a document from c1", correspondent=self.c1, checksum="A", document_type=self.dt, storage_path=self.sp1, ) self.doc2 = Document.objects.create( title="doc1", content="this is another document, but from c2", correspondent=self.c2, checksum="B", ) self.doc_inbox = Document.objects.create( title="doc235", content="aa", checksum="C", ) self.doc1.tags.add(self.t1) self.doc2.tags.add(self.t1) self.doc2.tags.add(self.t3) self.doc_inbox.tags.add(self.t2) def generate_train_and_save(self) -> None: """ Generates the training data, trains and saves the updated pickle file. This ensures the test is using the same scikit learn version and eliminates a warning from the test suite """ self.generate_test_data() self.classifier.train() self.classifier.save() def test_no_training_data(self) -> None: """ GIVEN: - No documents exist to train WHEN: - Classifier training is requested THEN: - Exception is raised """ with self.assertRaisesMessage(ValueError, "No training data available."): self.classifier.train() def test_no_non_inbox_tags(self) -> None: """ GIVEN: - No documents without an inbox tag exist WHEN: - Classifier training is requested THEN: - Exception is raised """ t1 = Tag.objects.create( name="t1", matching_algorithm=Tag.MATCH_ANY, pk=34, is_inbox_tag=True, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", ) doc1.tags.add(t1) with self.assertRaisesMessage(ValueError, "No training data available."): self.classifier.train() def testEmpty(self) -> None: """ GIVEN: - A document exists - No tags/not enough data to predict WHEN: - Classifier prediction is requested THEN: - Classifier returns no predictions """ Document.objects.create(title="WOW", checksum="3457", content="ASD") self.classifier.train() self.assertIsNone(self.classifier.document_type_classifier) self.assertIsNone(self.classifier.tags_classifier) self.assertIsNone(self.classifier.correspondent_classifier) self.assertListEqual(self.classifier.predict_tags(""), []) self.assertIsNone(self.classifier.predict_document_type("")) self.assertIsNone(self.classifier.predict_correspondent("")) def testTrain(self) -> None: """ GIVEN: - Test data WHEN: - Classifier is trained THEN: - Classifier uses correct values for correspondent learning - Classifier uses correct values for tags learning """ self.generate_test_data() self.classifier.train() self.assertListEqual( list(self.classifier.correspondent_classifier.classes_), [-1, self.c1.pk], ) self.assertListEqual( list(self.classifier.tags_binarizer.classes_), [self.t1.pk, self.t3.pk], ) def testPredict(self) -> None: """ GIVEN: - Classifier trained against test data WHEN: - Prediction requested for correspondent, tags, type THEN: - Expected predictions based on training set """ self.generate_test_data() self.classifier.train() with ( mock.patch.object( self.classifier.data_vectorizer, "transform", wraps=self.classifier.data_vectorizer.transform, ) as mock_transform, mock.patch.object( self.classifier, "preprocess_content", wraps=self.classifier.preprocess_content, ) as mock_preprocess_content, ): self.assertEqual( self.classifier.predict_correspondent(self.doc1.content), self.c1.pk, ) self.assertEqual( self.classifier.predict_correspondent(self.doc2.content), None, ) self.assertListEqual( self.classifier.predict_tags(self.doc1.content), [self.t1.pk], ) self.assertListEqual( self.classifier.predict_tags(self.doc2.content), [self.t1.pk, self.t3.pk], ) self.assertEqual( self.classifier.predict_document_type(self.doc1.content), self.dt.pk, ) self.assertEqual( self.classifier.predict_document_type(self.doc2.content), None, ) # Check that the classifier vectorized content and text preprocessing has been cached # It should be called once per document (doc1 and doc2) self.assertEqual(mock_preprocess_content.call_count, 2) self.assertEqual(mock_transform.call_count, 2) def test_vectorize_recomputes_tampered_cache_entry(self) -> None: cached = bytearray(signed_pickle_dumps(["cached vector"])) cached[HMAC_SIZE] ^= 0xFF self.classifier.data_vectorizer = mock.Mock() self.classifier.data_vectorizer.transform.return_value = ["fresh vector"] with ( mock.patch( "documents.classifier.read_cache.get", return_value=bytes(cached), ), mock.patch("documents.classifier.read_cache.set") as cache_set, mock.patch("documents.classifier.read_cache.touch") as cache_touch, ): result = self.classifier._vectorize("content") self.assertEqual(result, ["fresh vector"]) self.classifier.data_vectorizer.transform.assert_called_once() cache_set.assert_called_once() cache_touch.assert_not_called() def test_no_retrain_if_no_change(self) -> None: """ GIVEN: - Classifier trained with current data WHEN: - Classifier training is requested again THEN: - Classifier does not redo training """ self.generate_test_data() self.assertTrue(self.classifier.train()) self.assertFalse(self.classifier.train()) def test_retrain_if_change(self) -> None: """ GIVEN: - Classifier trained with current data WHEN: - Classifier training is requested again - Documents have changed THEN: - Classifier does not redo training """ self.generate_test_data() self.assertTrue(self.classifier.train()) self.doc1.correspondent = self.c2 self.doc1.save() self.assertTrue(self.classifier.train()) def test_retrain_if_auto_match_set_changed(self) -> None: """ GIVEN: - Classifier trained with current data WHEN: - Classifier training is requested again - Some new AUTO match object exists THEN: - Classifier does redo training """ self.generate_test_data() # Add the ANY type self.doc1.tags.add(self.t4) self.assertTrue(self.classifier.train()) # Change the matching type self.t4.matching_algorithm = MatchingModel.MATCH_AUTO self.t4.save() self.assertTrue(self.classifier.train()) def testVersionIncreased(self) -> None: """ GIVEN: - Existing classifier model saved at a version WHEN: - Attempt to load classifier file from newer version THEN: - Exception is raised """ self.generate_train_and_save() classifier2 = DocumentClassifier() current_ver = DocumentClassifier.FORMAT_VERSION with mock.patch( "documents.classifier.DocumentClassifier.FORMAT_VERSION", current_ver + 1, ): # assure that we won't load old classifiers. self.assertRaises(IncompatibleClassifierVersionError, classifier2.load) self.classifier.save() # assure that we can load the classifier after saving it. classifier2.load() def testSaveClassifier(self) -> None: self.generate_train_and_save() new_classifier = DocumentClassifier() new_classifier.load() new_classifier.preprocess_content = mock.MagicMock(side_effect=dummy_preprocess) self.assertFalse(new_classifier.train()) def test_load_and_classify(self) -> None: self.generate_train_and_save() new_classifier = DocumentClassifier() new_classifier.load() new_classifier.preprocess_content = mock.MagicMock(side_effect=dummy_preprocess) self.assertCountEqual(new_classifier.predict_tags(self.doc2.content), [45, 12]) def test_load_corrupt_file(self) -> None: """ GIVEN: - Corrupted classifier pickle file WHEN: - An attempt is made to load the classifier THEN: - The ClassifierModelCorruptError is raised """ self.generate_train_and_save() # Write garbage data (valid HMAC length but invalid content) Path(settings.MODEL_FILE).write_bytes(b"\x00" * 64) with self.assertRaises(ClassifierModelCorruptError): self.classifier.load() self.assertIsNone(load_classifier()) def test_load_corrupt_pickle_valid_hmac(self) -> None: """ GIVEN: - A classifier file with valid HMAC but unparsable pickle data WHEN: - An attempt is made to load the classifier THEN: - The ClassifierModelCorruptError is raised """ garbage_data = b"this is not valid pickle data" signature = DocumentClassifier._compute_hmac(garbage_data) Path(settings.MODEL_FILE).write_bytes(signature + garbage_data) with self.assertRaises(ClassifierModelCorruptError): self.classifier.load() def test_load_tampered_file(self) -> None: """ GIVEN: - A classifier model file whose data has been modified WHEN: - An attempt is made to load the classifier THEN: - The ClassifierModelCorruptError is raised due to HMAC mismatch """ self.generate_train_and_save() raw = Path(settings.MODEL_FILE).read_bytes() # Flip a byte in the data portion (after the 32-byte HMAC) tampered = raw[:32] + bytes([raw[32] ^ 0xFF]) + raw[33:] Path(settings.MODEL_FILE).write_bytes(tampered) with self.assertRaises(ClassifierModelCorruptError): self.classifier.load() def test_load_wrong_secret_key(self) -> None: """ GIVEN: - A classifier model file signed with a different SECRET_KEY WHEN: - An attempt is made to load the classifier THEN: - The ClassifierModelCorruptError is raised due to HMAC mismatch """ self.generate_train_and_save() with override_settings(SECRET_KEY="different-secret-key"): with self.assertRaises(ClassifierModelCorruptError): self.classifier.load() def test_load_truncated_file(self) -> None: """ GIVEN: - A classifier model file that is too short to contain an HMAC WHEN: - An attempt is made to load the classifier THEN: - The ClassifierModelCorruptError is raised """ Path(settings.MODEL_FILE).write_bytes(b"\x00" * 16) with self.assertRaises(ClassifierModelCorruptError): self.classifier.load() def test_load_new_scikit_learn_version(self) -> None: """ GIVEN: - classifier pickle file triggers an InconsistentVersionWarning WHEN: - An attempt is made to load the classifier THEN: - IncompatibleClassifierVersionError is raised """ from sklearn.exceptions import InconsistentVersionWarning self.generate_train_and_save() fake_warning = warnings.WarningMessage( message=InconsistentVersionWarning( estimator_name="MLPClassifier", current_sklearn_version="1.0", original_sklearn_version="0.9", ), category=InconsistentVersionWarning, filename="", lineno=0, ) real_catch_warnings = warnings.catch_warnings class PatchedCatchWarnings(real_catch_warnings): def __enter__(self): w = super().__enter__() w.append(fake_warning) return w with mock.patch( "documents.classifier.warnings.catch_warnings", PatchedCatchWarnings, ): with self.assertRaises(IncompatibleClassifierVersionError): self.classifier.load() def test_one_correspondent_predict(self) -> None: c1 = Correspondent.objects.create( name="c1", matching_algorithm=Correspondent.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", correspondent=c1, checksum="A", ) self.classifier.train() self.assertEqual(self.classifier.predict_correspondent(doc1.content), c1.pk) def test_one_correspondent_predict_manydocs(self) -> None: c1 = Correspondent.objects.create( name="c1", matching_algorithm=Correspondent.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", correspondent=c1, checksum="A", ) doc2 = Document.objects.create( title="doc2", content="this is a document from no one", checksum="B", ) self.classifier.train() self.assertEqual(self.classifier.predict_correspondent(doc1.content), c1.pk) self.assertIsNone(self.classifier.predict_correspondent(doc2.content)) def test_one_type_predict(self) -> None: dt = DocumentType.objects.create( name="dt", matching_algorithm=DocumentType.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", document_type=dt, ) self.classifier.train() self.assertEqual(self.classifier.predict_document_type(doc1.content), dt.pk) def test_one_type_predict_manydocs(self) -> None: dt = DocumentType.objects.create( name="dt", matching_algorithm=DocumentType.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", document_type=dt, ) doc2 = Document.objects.create( title="doc1", content="this is a document from c2", checksum="B", ) self.classifier.train() self.assertEqual(self.classifier.predict_document_type(doc1.content), dt.pk) self.assertIsNone(self.classifier.predict_document_type(doc2.content)) def test_one_path_predict(self) -> None: sp = StoragePath.objects.create( name="sp", matching_algorithm=StoragePath.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", storage_path=sp, ) self.classifier.train() self.assertEqual(self.classifier.predict_storage_path(doc1.content), sp.pk) def test_one_path_predict_manydocs(self) -> None: sp = StoragePath.objects.create( name="sp", matching_algorithm=StoragePath.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", storage_path=sp, ) doc2 = Document.objects.create( title="doc1", content="this is a document from c2", checksum="B", ) self.classifier.train() self.assertEqual(self.classifier.predict_storage_path(doc1.content), sp.pk) self.assertIsNone(self.classifier.predict_storage_path(doc2.content)) def test_predict_rejects_prediction_below_match_threshold(self) -> None: """ GIVEN: - Classifiers trained against test data with confident predictions WHEN: - CLASSIFIER_MATCH_THRESHOLD exceeds the model's confidence THEN: - Every predict_* method discards the match in favor of no match """ c1 = Correspondent.objects.create( name="c1", matching_algorithm=Correspondent.MATCH_AUTO, ) dt1 = DocumentType.objects.create( name="dt1", matching_algorithm=DocumentType.MATCH_AUTO, ) sp1 = StoragePath.objects.create( name="sp1", matching_algorithm=StoragePath.MATCH_AUTO, ) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", correspondent=c1, document_type=dt1, storage_path=sp1, checksum="A", ) Document.objects.create( title="doc2", content="this is a document from no one", checksum="B", ) self.classifier.train() predictors = { "correspondent": self.classifier.predict_correspondent, "document_type": self.classifier.predict_document_type, "storage_path": self.classifier.predict_storage_path, } # No real prediction can reach a confidence this high, so this # isolates the threshold check from the model's actual output. with override_settings(CLASSIFIER_MATCH_THRESHOLD=0.999999): for name, predict in predictors.items(): with self.subTest(field=name): self.assertIsNone(predict(doc1.content)) def test_train_uses_balanced_sample_weight(self) -> None: """ GIVEN: - A training set with correspondents, document types and storage paths WHEN: - The classifier is trained THEN: - Each MLP classifier is fit with balanced sample weights, so that over-represented classes don't dominate predictions """ c1 = Correspondent.objects.create( name="c1", matching_algorithm=Correspondent.MATCH_AUTO, ) dt1 = DocumentType.objects.create( name="dt1", matching_algorithm=DocumentType.MATCH_AUTO, ) sp1 = StoragePath.objects.create( name="sp1", matching_algorithm=StoragePath.MATCH_AUTO, ) Document.objects.create( title="doc1", content="this is a document from c1", correspondent=c1, document_type=dt1, storage_path=sp1, checksum="A", ) Document.objects.create( title="doc2", content="this is a document from no one", checksum="B", ) with mock.patch( "sklearn.utils.class_weight.compute_sample_weight", return_value=None, ) as mocked_compute_sample_weight: self.classifier.train() self.assertEqual(mocked_compute_sample_weight.call_count, 3) for call in mocked_compute_sample_weight.call_args_list: self.assertEqual(call.args[0], "balanced") def test_one_tag_predict(self) -> None: t1 = Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", ) doc1.tags.add(t1) self.classifier.train() self.assertListEqual(self.classifier.predict_tags(doc1.content), [t1.pk]) def test_one_tag_predict_unassigned(self) -> None: Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", ) self.classifier.train() self.assertListEqual(self.classifier.predict_tags(doc1.content), []) def test_two_tags_predict_singledoc(self) -> None: t1 = Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12) t2 = Tag.objects.create(name="t2", matching_algorithm=Tag.MATCH_AUTO, pk=121) doc4 = Document.objects.create( title="doc1", content="this is a document from c4", checksum="D", ) doc4.tags.add(t1) doc4.tags.add(t2) self.classifier.train() self.assertListEqual(self.classifier.predict_tags(doc4.content), [t1.pk, t2.pk]) def test_two_tags_predict(self) -> None: t1 = Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12) t2 = Tag.objects.create(name="t2", matching_algorithm=Tag.MATCH_AUTO, pk=121) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", ) doc2 = Document.objects.create( title="doc1", content="this is a document from c2", checksum="B", ) doc3 = Document.objects.create( title="doc1", content="this is a document from c3", checksum="C", ) doc4 = Document.objects.create( title="doc1", content="this is a document from c4", checksum="D", ) doc1.tags.add(t1) doc2.tags.add(t2) doc4.tags.add(t1) doc4.tags.add(t2) self.classifier.train() self.assertListEqual(self.classifier.predict_tags(doc1.content), [t1.pk]) self.assertListEqual(self.classifier.predict_tags(doc2.content), [t2.pk]) self.assertListEqual(self.classifier.predict_tags(doc3.content), []) self.assertListEqual(self.classifier.predict_tags(doc4.content), [t1.pk, t2.pk]) def test_one_tag_predict_multi(self) -> None: t1 = Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", ) doc2 = Document.objects.create( title="doc2", content="this is a document from c2", checksum="B", ) doc1.tags.add(t1) doc2.tags.add(t1) self.classifier.train() self.assertListEqual(self.classifier.predict_tags(doc1.content), [t1.pk]) self.assertListEqual(self.classifier.predict_tags(doc2.content), [t1.pk]) def test_one_tag_predict_multi_2(self) -> None: t1 = Tag.objects.create(name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12) doc1 = Document.objects.create( title="doc1", content="this is a document from c1", checksum="A", ) doc2 = Document.objects.create( title="doc2", content="this is a document from c2", checksum="B", ) doc1.tags.add(t1) self.classifier.train() self.assertListEqual(self.classifier.predict_tags(doc1.content), [t1.pk]) self.assertListEqual(self.classifier.predict_tags(doc2.content), []) def test_load_classifier_not_exists(self) -> None: self.assertFalse(Path(settings.MODEL_FILE).exists()) self.assertIsNone(load_classifier()) @mock.patch("documents.classifier.DocumentClassifier.load") def test_load_classifier(self, load) -> None: Path(settings.MODEL_FILE).touch() self.assertIsNotNone(load_classifier()) load.assert_called_once() @override_settings( CACHES={ "default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"}, }, ) @override_settings( MODEL_FILE=str(Path(__file__).parent / "data" / "model.pickle"), ) @pytest.mark.skip( reason="Disabled caching due to high memory usage - need to investigate.", ) def test_load_classifier_cached(self) -> None: classifier = load_classifier() self.assertIsNotNone(classifier) with mock.patch("documents.classifier.DocumentClassifier.load") as load: load_classifier() load.assert_not_called() @mock.patch("documents.classifier.DocumentClassifier.load") def test_load_classifier_incompatible_version(self, load) -> None: Path(settings.MODEL_FILE).touch() self.assertTrue(Path(settings.MODEL_FILE).exists()) load.side_effect = IncompatibleClassifierVersionError("Dummy Error") self.assertIsNone(load_classifier()) self.assertFalse(Path(settings.MODEL_FILE).exists()) @mock.patch("documents.classifier.DocumentClassifier.load") def test_load_classifier_os_error(self, load) -> None: Path(settings.MODEL_FILE).touch() self.assertTrue(Path(settings.MODEL_FILE).exists()) load.side_effect = OSError() self.assertIsNone(load_classifier()) self.assertTrue(Path(settings.MODEL_FILE).exists()) @mock.patch("documents.classifier.DocumentClassifier.load") def test_load_classifier_raise_exception(self, mock_load) -> None: Path(settings.MODEL_FILE).touch() mock_load.side_effect = IncompatibleClassifierVersionError("Dummy Error") with self.assertRaises(IncompatibleClassifierVersionError): load_classifier(raise_exception=True) Path(settings.MODEL_FILE).touch() mock_load.side_effect = ClassifierModelCorruptError() with self.assertRaises(ClassifierModelCorruptError): load_classifier(raise_exception=True) Path(settings.MODEL_FILE).touch() mock_load.side_effect = OSError() with self.assertRaises(OSError): load_classifier(raise_exception=True) Path(settings.MODEL_FILE).touch() mock_load.side_effect = Exception() with self.assertRaises(Exception): load_classifier(raise_exception=True) class TestClassifierSave: @pytest.fixture def model_file(self, tmp_path: Path, settings: Settings) -> Path: settings.MODEL_FILE = tmp_path / "classifier.pickle" return settings.MODEL_FILE @pytest.fixture def classifier(self) -> DocumentClassifier: classifier = DocumentClassifier() classifier.last_doc_change_time = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) classifier.last_auto_type_hash = b"\x01" * 32 return classifier def test_save_writes_signed_pickle( self, model_file: Path, classifier: DocumentClassifier, ) -> None: """ GIVEN: - A classifier with training state WHEN: - The classifier is saved THEN: - The file is the HMAC of the pickled state followed by that pickle - The pickle uses the highest protocol - The saved state loads back into a new classifier - No temporary file is left behind """ classifier.save() raw = model_file.read_bytes() signature = raw[: DocumentClassifier.HMAC_SIZE] data = raw[DocumentClassifier.HMAC_SIZE :] assert signature == DocumentClassifier._compute_hmac(data) # A pickle opens with the PROTO opcode followed by the protocol number assert data[:2] == bytes([pickle.PROTO[0], pickle.HIGHEST_PROTOCOL]) assert pickle.loads(data)[:3] == ( DocumentClassifier.FORMAT_VERSION, classifier.last_doc_change_time, classifier.last_auto_type_hash, ) loaded = DocumentClassifier() loaded.load() assert loaded.last_doc_change_time == classifier.last_doc_change_time assert loaded.last_auto_type_hash == classifier.last_auto_type_hash assert not model_file.with_name(f"{model_file.name}.part").exists() def test_save_failure_removes_partial_file( self, model_file: Path, classifier: DocumentClassifier, mocker: MockerFixture, ) -> None: """ GIVEN: - An existing classifier model file WHEN: - Saving a new classifier fails part way through writing THEN: - The error is raised - The partially written temporary file is removed - The existing model file is left untouched """ model_file.write_bytes(b"existing model") mocker.patch( "documents.classifier.pickle.dump", side_effect=RuntimeError("disk full"), ) with pytest.raises(RuntimeError, match="disk full"): classifier.save() assert not model_file.with_name(f"{model_file.name}.part").exists() assert model_file.read_bytes() == b"existing model" class _StubProbaClassifier: """ A fake scikit-learn classifier exposing just enough of the API for `_predict_with_threshold`: `classes_` and `predict_proba`. """ def __init__(self, classes: list[int], probabilities: list[float]) -> None: self.classes_ = np.array(classes) self._probabilities = np.array([probabilities]) def predict_proba(self, X) -> np.ndarray: return self._probabilities @pytest.mark.parametrize( ("classes", "probabilities", "threshold", "expected"), [ # confident prediction above the threshold is returned ([-1, 3], [0.1, 0.9], 0.6, 3), # prediction below the threshold is discarded ([-1, 3], [0.45, 0.55], 0.6, None), # boundary: exactly at the threshold is accepted, not discarded ([-1, 3], [0.4, 0.6], 0.6, 3), # the winning class is the "no match" pseudo-class, regardless of its # own confidence ([-1, 3], [0.99, 0.01], 0.0, None), # threshold of 0.0 disables the confidence check entirely ([-1, 3], [0.45, 0.55], 0.0, 3), ], ) def test_predict_with_threshold(classes, probabilities, threshold, expected) -> None: classifier = _StubProbaClassifier(classes, probabilities) result = _predict_with_threshold(classifier, X=None, threshold=threshold) assert result == expected def test_classifier_match_threshold_default() -> None: """ GIVEN: - No PAPERLESS_CLASSIFIER_MATCH_THRESHOLD environment variable is set THEN: - The classifier match threshold defaults to 0.3 """ assert settings.CLASSIFIER_MATCH_THRESHOLD == 0.3 class TestPreprocessContent: @pytest.fixture def samples(self) -> Path: return Path(__file__).parent / "samples" @pytest.fixture def content(self, samples: Path) -> str: return (samples / "content.txt").read_text() def test_supported_language( self, settings: Settings, samples: Path, content: str, ) -> None: """ GIVEN: - The classifier language is English, the default WHEN: - Document content is preprocessed THEN: - Stop words are removed and the remaining words are stemmed """ settings.CLASSIFIER_LANGUAGE = "english" expected = (samples / "preprocessed_content_advanced.txt").read_text() assert DocumentClassifier().preprocess_content(content) == expected.rstrip() def test_unsupported_language( self, settings: Settings, samples: Path, content: str, ) -> None: """ GIVEN: - No classifier language (the OCR language has no stemming support) WHEN: - Document content is preprocessed THEN: - The content is only lowercased and split into words """ settings.CLASSIFIER_LANGUAGE = None expected = (samples / "preprocessed_content.txt").read_text() assert DocumentClassifier().preprocess_content(content) == expected.rstrip() @pytest.mark.parametrize( "language", [pytest.param("english", id="supported"), pytest.param(None, id="unsupported")], ) def test_empty_content(self, settings: Settings, language: str | None) -> None: """ GIVEN: - Empty document content WHEN: - The content is preprocessed THEN: - The result is empty """ settings.CLASSIFIER_LANGUAGE = language assert DocumentClassifier().preprocess_content("") == "" @pytest.mark.django_db class TestClassifierTrainTagLabels: @pytest.fixture(autouse=True) def _simple_preprocess(self, mocker: MockerFixture) -> None: mocker.patch.object( DocumentClassifier, "preprocess_content", side_effect=dummy_preprocess, ) @pytest.fixture def auto_tags(self) -> list[Tag]: return TagFactory.create_batch(2, matching_algorithm=MatchingModel.MATCH_AUTO) def test_train_query_count_does_not_scale_with_documents( self, auto_tags: list[Tag], ) -> None: """ GIVEN: - Documents with auto matching tags WHEN: - The classifier is trained, then more documents are added and it is trained again THEN: - Both trainings run the same number of queries """ for doc in DocumentFactory.create_batch(2): doc.tags.set(auto_tags) with CaptureQueriesContext(connection) as few_documents: DocumentClassifier().train() for doc in DocumentFactory.create_batch(6): doc.tags.set(auto_tags) with CaptureQueriesContext(connection) as more_documents: DocumentClassifier().train() assert len(more_documents) == len(few_documents) def test_train_uses_only_auto_tags_as_labels( self, auto_tags: list[Tag], ) -> None: """ GIVEN: - Documents with both auto matching and non auto matching tags WHEN: - The classifier is trained THEN: - Only the auto matching tags are used as tag labels """ manual_tag = TagFactory(matching_algorithm=MatchingModel.MATCH_ANY) first, second, third = DocumentFactory.create_batch(3) first.tags.set([auto_tags[0], manual_tag]) second.tags.set([auto_tags[1], manual_tag]) third.tags.set([*auto_tags, manual_tag]) classifier = DocumentClassifier() classifier.train() assert list(classifier.tags_binarizer.classes_) == sorted( tag.pk for tag in auto_tags ) @pytest.mark.django_db class TestClassifierTrainContent: def test_train_content_follows_label_order_across_chunks( self, mocker: MockerFixture, ) -> None: """ GIVEN: - More documents than fit in one content chunk WHEN: - The classifier is trained THEN: - Every document's content is preprocessed once, in document order """ mocker.patch("documents.classifier._CONTENT_CHUNK_SIZE", 2) docs = DocumentFactory.create_batch(5) preprocess = mocker.patch.object( DocumentClassifier, "preprocess_content", side_effect=dummy_preprocess, ) DocumentClassifier().train() assert [call.args[0] for call in preprocess.call_args_list] == [ doc.content for doc in sorted(docs, key=lambda doc: doc.pk) ] def test_train_document_deleted_while_training( self, mocker: MockerFixture, ) -> None: """ GIVEN: - Two documents WHEN: - The second document is deleted after its labels were gathered, but before its content is fetched THEN: - Training completes - The deleted document is trained with empty content, keeping labels and content aligned """ mocker.patch("documents.classifier._CONTENT_CHUNK_SIZE", 1) first, second = DocumentFactory.create_batch(2) def delete_second_then_preprocess(content: str, **kwargs) -> str: if content == first.content: second.delete() return dummy_preprocess(content) preprocess = mocker.patch.object( DocumentClassifier, "preprocess_content", side_effect=delete_second_then_preprocess, ) assert DocumentClassifier().train() assert [call.args[0] for call in preprocess.call_args_list] == [ first.content, "", ] class TestTextAnalyzer: @pytest.mark.parametrize( "language", [ pytest.param(language, id=language) for language in sorted(set(CLASSIFIER_LANGUAGES.values())) ], ) def test_builds_for_every_classifier_language(self, language: str) -> None: """ GIVEN: - A language the classifier supports WHEN: - Text is analyzed with its classifier language THEN: - Tokens are produced """ assert _text_analyzer(language).analyze("Paperless invoice 2026") def test_english_removes_snowball_stop_words(self) -> None: """ GIVEN: - English text with a contraction and stop words missing from Tantivy's own English list WHEN: - The text is analyzed THEN: - All stop words are removed, including the contraction - The remaining words are stemmed """ tokens = _text_analyzer("english").analyze( "They were about to pay the invoices, don't worry", ) assert tokens == ["pay", "invoic", "worri"] def test_keeps_underscores_within_tokens(self) -> None: """ GIVEN: - Text with a word joined by an underscore WHEN: - The text is analyzed THEN: - The word stays one token """ tokens = _text_analyzer("english").analyze("tax_id") assert tokens == ["tax_id"]