paradizelost/paperless-ngx
1
0
Fork 0
mirror of https://github.com/paperless-ngx/paperless-ngx.git synced 2025-08-01 18:37:42 -05:00
Code Issues Packages Projects Releases Wiki Activity

Format Python code with black

Browse Source
This commit is contained in:
kpj
2022-02-27 15:26:41 +01:00
parent f0ffc69010
commit c56cb25b5f
136 changed files with 6142 additions and 3811 deletions
Download Patch File Download Diff File Expand all files Collapse all files

197
src/documents/tests/test_classifier.py
View File

@@ -7,30 +7,60 @@ import pytest
from django.conf import settings
from django.test import TestCase, override_settings
from documents.classifier import DocumentClassifier, IncompatibleClassifierVersionError, load_classifier
from documents.classifier import (
DocumentClassifier,
IncompatibleClassifierVersionError,
load_classifier,
)
from documents.models import Correspondent, Document, Tag, DocumentType
from documents.tests.utils import DirectoriesMixin
class TestClassifier(DirectoriesMixin, TestCase):
def setUp(self):
super(TestClassifier, self).setUp()
self.classifier = DocumentClassifier()
def generate_test_data(self):
self.c1 = Correspondent.objects.create(name="c1", matching_algorithm=Correspondent.MATCH_AUTO)
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.dt = DocumentType.objects.create(name="dt", matching_algorithm=DocumentType.MATCH_AUTO)
self.dt2 = DocumentType.objects.create(name="dt2", matching_algorithm=DocumentType.MATCH_AUTO)
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.dt = DocumentType.objects.create(
name="dt", matching_algorithm=DocumentType.MATCH_AUTO
)
self.dt2 = DocumentType.objects.create(
name="dt2", matching_algorithm=DocumentType.MATCH_AUTO
)
self.doc1 = Document.objects.create(title="doc1", content="this is a document from c1", correspondent=self.c1, checksum="A", document_type=self.dt)
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 = Document.objects.create(
title="doc1",
content="this is a document from c1",
correspondent=self.c1,
checksum="A",
document_type=self.dt,
)
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)
@@ -59,17 +89,29 @@ class TestClassifier(DirectoriesMixin, TestCase):
def testTrain(self):
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])
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):
self.generate_test_data()
self.classifier.train()
self.assertEqual(self.classifier.predict_correspondent(self.doc1.content), self.c1.pk)
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.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)
def testDatasetHashing(self):
@@ -90,7 +132,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
classifier2 = DocumentClassifier()
current_ver = DocumentClassifier.FORMAT_VERSION
with mock.patch("documents.classifier.DocumentClassifier.FORMAT_VERSION", current_ver+1):
with mock.patch(
"documents.classifier.DocumentClassifier.FORMAT_VERSION", current_ver + 1
):
# assure that we won't load old classifiers.
self.assertRaises(IncompatibleClassifierVersionError, classifier2.load)
@@ -112,7 +156,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
new_classifier.load()
self.assertFalse(new_classifier.train())
@override_settings(MODEL_FILE=os.path.join(os.path.dirname(__file__), "data", "model.pickle"))
@override_settings(
MODEL_FILE=os.path.join(os.path.dirname(__file__), "data", "model.pickle")
)
def test_load_and_classify(self):
self.generate_test_data()
@@ -122,38 +168,67 @@ class TestClassifier(DirectoriesMixin, TestCase):
self.assertCountEqual(new_classifier.predict_tags(self.doc2.content), [45, 12])
def test_one_correspondent_predict(self):
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")
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):
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 noone", checksum="B")
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 noone", 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):
dt = DocumentType.objects.create(name="dt", matching_algorithm=DocumentType.MATCH_AUTO)
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)
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):
dt = DocumentType.objects.create(name="dt", matching_algorithm=DocumentType.MATCH_AUTO)
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)
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")
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)
@@ -162,7 +237,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
def test_one_tag_predict(self):
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 = Document.objects.create(
title="doc1", content="this is a document from c1", checksum="A"
)
doc1.tags.add(t1)
self.classifier.train()
@@ -171,7 +248,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
def test_one_tag_predict_unassigned(self):
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 = 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), [])
@@ -180,7 +259,9 @@ class TestClassifier(DirectoriesMixin, TestCase):
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 = Document.objects.create(
title="doc1", content="this is a document from c4", checksum="D"
)
doc4.tags.add(t1)
doc4.tags.add(t2)
@@ -191,10 +272,18 @@ class TestClassifier(DirectoriesMixin, TestCase):
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 = 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)
@@ -210,8 +299,12 @@ class TestClassifier(DirectoriesMixin, TestCase):
def test_one_tag_predict_multi(self):
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 = 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)
@@ -222,8 +315,12 @@ class TestClassifier(DirectoriesMixin, TestCase):
def test_one_tag_predict_multi_2(self):
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 = 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()
@@ -240,9 +337,15 @@ class TestClassifier(DirectoriesMixin, TestCase):
self.assertIsNotNone(load_classifier())
load.assert_called_once()
@override_settings(CACHES={'default': {'BACKEND': 'django.core.cache.backends.locmem.LocMemCache'}})
@override_settings(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.")
@override_settings(
CACHES={"default": {"BACKEND": "django.core.cache.backends.locmem.LocMemCache"}}
)
@override_settings(
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."
)
def test_load_classifier_cached(self):
classifier = load_classifier()
self.assertIsNotNone(classifier)