mirror of
https://github.com/paperless-ngx/paperless-ngx.git
synced 2025-08-01 18:37:42 -05:00
Runs the pre-commit hooks over all the Python files
This commit is contained in:
@@ -5,14 +5,15 @@ from unittest import mock
|
||||
|
||||
import pytest
|
||||
from django.conf import settings
|
||||
from django.test import TestCase, override_settings
|
||||
|
||||
from documents.classifier import (
|
||||
DocumentClassifier,
|
||||
IncompatibleClassifierVersionError,
|
||||
load_classifier,
|
||||
)
|
||||
from documents.models import Correspondent, Document, Tag, DocumentType
|
||||
from django.test import override_settings
|
||||
from django.test import TestCase
|
||||
from documents.classifier import DocumentClassifier
|
||||
from documents.classifier import IncompatibleClassifierVersionError
|
||||
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 Tag
|
||||
from documents.tests.utils import DirectoriesMixin
|
||||
|
||||
|
||||
@@ -23,26 +24,37 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
|
||||
def generate_test_data(self):
|
||||
self.c1 = Correspondent.objects.create(
|
||||
name="c1", matching_algorithm=Correspondent.MATCH_AUTO
|
||||
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
|
||||
name="c3",
|
||||
matching_algorithm=Correspondent.MATCH_AUTO,
|
||||
)
|
||||
self.t1 = Tag.objects.create(
|
||||
name="t1", matching_algorithm=Tag.MATCH_AUTO, pk=12
|
||||
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
|
||||
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
|
||||
name="t3",
|
||||
matching_algorithm=Tag.MATCH_AUTO,
|
||||
pk=45,
|
||||
)
|
||||
self.dt = DocumentType.objects.create(
|
||||
name="dt", matching_algorithm=DocumentType.MATCH_AUTO
|
||||
name="dt",
|
||||
matching_algorithm=DocumentType.MATCH_AUTO,
|
||||
)
|
||||
self.dt2 = DocumentType.objects.create(
|
||||
name="dt2", matching_algorithm=DocumentType.MATCH_AUTO
|
||||
name="dt2",
|
||||
matching_algorithm=DocumentType.MATCH_AUTO,
|
||||
)
|
||||
|
||||
self.doc1 = Document.objects.create(
|
||||
@@ -59,7 +71,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
checksum="B",
|
||||
)
|
||||
self.doc_inbox = Document.objects.create(
|
||||
title="doc235", content="aa", checksum="C"
|
||||
title="doc235",
|
||||
content="aa",
|
||||
checksum="C",
|
||||
)
|
||||
|
||||
self.doc1.tags.add(self.t1)
|
||||
@@ -90,27 +104,33 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
self.generate_test_data()
|
||||
self.classifier.train()
|
||||
self.assertListEqual(
|
||||
list(self.classifier.correspondent_classifier.classes_), [-1, self.c1.pk]
|
||||
list(self.classifier.correspondent_classifier.classes_),
|
||||
[-1, self.c1.pk],
|
||||
)
|
||||
self.assertListEqual(
|
||||
list(self.classifier.tags_binarizer.classes_), [self.t1.pk, self.t3.pk]
|
||||
list(self.classifier.tags_binarizer.classes_),
|
||||
[self.t1.pk, self.t3.pk],
|
||||
)
|
||||
|
||||
def testPredict(self):
|
||||
self.generate_test_data()
|
||||
self.classifier.train()
|
||||
self.assertEqual(
|
||||
self.classifier.predict_correspondent(self.doc1.content), self.c1.pk
|
||||
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.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.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.classifier.predict_document_type(self.doc1.content),
|
||||
self.dt.pk,
|
||||
)
|
||||
self.assertEqual(self.classifier.predict_document_type(self.doc2.content), None)
|
||||
|
||||
@@ -133,7 +153,8 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
|
||||
current_ver = DocumentClassifier.FORMAT_VERSION
|
||||
with mock.patch(
|
||||
"documents.classifier.DocumentClassifier.FORMAT_VERSION", current_ver + 1
|
||||
"documents.classifier.DocumentClassifier.FORMAT_VERSION",
|
||||
current_ver + 1,
|
||||
):
|
||||
# assure that we won't load old classifiers.
|
||||
self.assertRaises(IncompatibleClassifierVersionError, classifier2.load)
|
||||
@@ -157,7 +178,7 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
self.assertFalse(new_classifier.train())
|
||||
|
||||
@override_settings(
|
||||
MODEL_FILE=os.path.join(os.path.dirname(__file__), "data", "model.pickle")
|
||||
MODEL_FILE=os.path.join(os.path.dirname(__file__), "data", "model.pickle"),
|
||||
)
|
||||
def test_load_and_classify(self):
|
||||
self.generate_test_data()
|
||||
@@ -169,7 +190,8 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
|
||||
def test_one_correspondent_predict(self):
|
||||
c1 = Correspondent.objects.create(
|
||||
name="c1", matching_algorithm=Correspondent.MATCH_AUTO
|
||||
name="c1",
|
||||
matching_algorithm=Correspondent.MATCH_AUTO,
|
||||
)
|
||||
doc1 = Document.objects.create(
|
||||
title="doc1",
|
||||
@@ -183,7 +205,8 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
|
||||
def test_one_correspondent_predict_manydocs(self):
|
||||
c1 = Correspondent.objects.create(
|
||||
name="c1", matching_algorithm=Correspondent.MATCH_AUTO
|
||||
name="c1",
|
||||
matching_algorithm=Correspondent.MATCH_AUTO,
|
||||
)
|
||||
doc1 = Document.objects.create(
|
||||
title="doc1",
|
||||
@@ -192,7 +215,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
checksum="A",
|
||||
)
|
||||
doc2 = Document.objects.create(
|
||||
title="doc2", content="this is a document from noone", checksum="B"
|
||||
title="doc2",
|
||||
content="this is a document from noone",
|
||||
checksum="B",
|
||||
)
|
||||
|
||||
self.classifier.train()
|
||||
@@ -201,7 +226,8 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
|
||||
def test_one_type_predict(self):
|
||||
dt = DocumentType.objects.create(
|
||||
name="dt", matching_algorithm=DocumentType.MATCH_AUTO
|
||||
name="dt",
|
||||
matching_algorithm=DocumentType.MATCH_AUTO,
|
||||
)
|
||||
|
||||
doc1 = Document.objects.create(
|
||||
@@ -216,7 +242,8 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
|
||||
def test_one_type_predict_manydocs(self):
|
||||
dt = DocumentType.objects.create(
|
||||
name="dt", matching_algorithm=DocumentType.MATCH_AUTO
|
||||
name="dt",
|
||||
matching_algorithm=DocumentType.MATCH_AUTO,
|
||||
)
|
||||
|
||||
doc1 = Document.objects.create(
|
||||
@@ -227,7 +254,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
)
|
||||
|
||||
doc2 = Document.objects.create(
|
||||
title="doc1", content="this is a document from c2", checksum="B"
|
||||
title="doc1",
|
||||
content="this is a document from c2",
|
||||
checksum="B",
|
||||
)
|
||||
|
||||
self.classifier.train()
|
||||
@@ -238,7 +267,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
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"
|
||||
title="doc1",
|
||||
content="this is a document from c1",
|
||||
checksum="A",
|
||||
)
|
||||
|
||||
doc1.tags.add(t1)
|
||||
@@ -249,7 +280,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
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"
|
||||
title="doc1",
|
||||
content="this is a document from c1",
|
||||
checksum="A",
|
||||
)
|
||||
|
||||
self.classifier.train()
|
||||
@@ -260,7 +293,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
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"
|
||||
title="doc1",
|
||||
content="this is a document from c4",
|
||||
checksum="D",
|
||||
)
|
||||
|
||||
doc4.tags.add(t1)
|
||||
@@ -273,16 +308,24 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
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"
|
||||
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"
|
||||
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"
|
||||
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"
|
||||
title="doc1",
|
||||
content="this is a document from c4",
|
||||
checksum="D",
|
||||
)
|
||||
|
||||
doc1.tags.add(t1)
|
||||
@@ -300,10 +343,14 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
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"
|
||||
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"
|
||||
title="doc2",
|
||||
content="this is a document from c2",
|
||||
checksum="B",
|
||||
)
|
||||
|
||||
doc1.tags.add(t1)
|
||||
@@ -316,10 +363,14 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
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"
|
||||
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"
|
||||
title="doc2",
|
||||
content="this is a document from c2",
|
||||
checksum="B",
|
||||
)
|
||||
|
||||
doc1.tags.add(t1)
|
||||
@@ -338,13 +389,15 @@ class TestClassifier(DirectoriesMixin, TestCase):
|
||||
load.assert_called_once()
|
||||
|
||||
@override_settings(
|
||||
CACHES={"default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"}}
|
||||
CACHES={
|
||||
"default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"},
|
||||
},
|
||||
)
|
||||
@override_settings(
|
||||
MODEL_FILE=os.path.join(os.path.dirname(__file__), "data", "model.pickle")
|
||||
MODEL_FILE=os.path.join(os.path.dirname(__file__), "data", "model.pickle"),
|
||||
)
|
||||
@pytest.mark.skip(
|
||||
reason="Disabled caching due to high memory usage - need to investigate."
|
||||
reason="Disabled caching due to high memory usage - need to investigate.",
|
||||
)
|
||||
def test_load_classifier_cached(self):
|
||||
classifier = load_classifier()
|
||||
|
Reference in New Issue
Block a user