From 4f391e47114bf1777fc030c7b15a636732f68c4e Mon Sep 17 00:00:00 2001
From: Trenton Holmes <797416+stumpylog@users.noreply.github.com>
Date: Wed, 20 Sep 2023 17:04:01 -0700
Subject: [PATCH] A couple small improvements and a fix for the mail source
 overriding

---
 src/documents/data_models.py | 10 +++++-----
 src/documents/models.py      | 20 +++++++++-----------
 src/documents/serialisers.py |  9 ++++-----
 src/documents/views.py       | 10 ++++++----
 4 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/src/documents/data_models.py b/src/documents/data_models.py
index a2188f4d6..99cf92457 100644
--- a/src/documents/data_models.py
+++ b/src/documents/data_models.py
@@ -1,6 +1,6 @@
 import dataclasses
 import datetime
-import enum
+from enum import IntEnum
 from pathlib import Path
 from typing import Optional
 
@@ -30,14 +30,14 @@ class DocumentMetadataOverrides:
     change_groups: Optional[list[int]] = None
 
 
-class DocumentSource(enum.IntEnum):
+class DocumentSource(IntEnum):
     """
     The source of an incoming document.  May have other uses in the future
     """
 
-    ConsumeFolder = enum.auto()
-    ApiUpload = enum.auto()
-    MailFetch = enum.auto()
+    ConsumeFolder = 1
+    ApiUpload = 2
+    MailFetch = 3
 
 
 @dataclasses.dataclass
diff --git a/src/documents/models.py b/src/documents/models.py
index 60ee42dcf..8f3c9a0de 100644
--- a/src/documents/models.py
+++ b/src/documents/models.py
@@ -23,9 +23,6 @@ from multiselectfield import MultiSelectField
 from documents.data_models import DocumentSource
 from documents.parsers import get_default_file_extension
 
-ALL_STATES = sorted(states.ALL_STATES)
-TASK_STATE_CHOICES = sorted(zip(ALL_STATES, ALL_STATES))
-
 
 class ModelWithOwner(models.Model):
     owner = models.ForeignKey(
@@ -575,6 +572,9 @@ class UiSettings(models.Model):
 
 
 class PaperlessTask(models.Model):
+    ALL_STATES = sorted(states.ALL_STATES)
+    TASK_STATE_CHOICES = sorted(zip(ALL_STATES, ALL_STATES))
+
     task_id = models.CharField(
         max_length=255,
         unique=True,
@@ -740,14 +740,12 @@ class ShareLink(models.Model):
         return f"Share Link for {self.document.title}"
 
 
-DOCUMENT_SOURCE = (
-    (int(DocumentSource.ConsumeFolder), _("Consume Folder")),
-    (int(DocumentSource.ApiUpload), _("Api Upload")),
-    (int(DocumentSource.MailFetch), _("Mail Fetch")),
-)
-
-
 class ConsumptionTemplate(ModelWithOwner):
+    class DocumentSourceChoices(models.IntegerChoices):
+        CONSUME_FOLDER = DocumentSource.ConsumeFolder.value, _("Consume Folder")
+        API_UPLOAD = DocumentSource.ApiUpload.value, _("Api Upload")
+        MAIL_FETCH = DocumentSource.MailFetch.value, _("Mail Fetch")
+
     class Meta:
         verbose_name = _("consumption template")
         verbose_name_plural = _("consumption templates")
@@ -758,7 +756,7 @@ class ConsumptionTemplate(ModelWithOwner):
 
     sources = MultiSelectField(
         max_length=3,
-        choices=DOCUMENT_SOURCE,
+        choices=DocumentSourceChoices.choices,
         default=f"{DocumentSource.ConsumeFolder},{DocumentSource.ApiUpload},{DocumentSource.MailFetch}",
     )
 
diff --git a/src/documents/serialisers.py b/src/documents/serialisers.py
index 250099f11..ab957c1c9 100644
--- a/src/documents/serialisers.py
+++ b/src/documents/serialisers.py
@@ -18,12 +18,11 @@ from rest_framework import serializers
 from rest_framework.fields import SerializerMethodField
 
 from documents.data_models import DocumentSource
+from documents.models import ConsumptionTemplate
 from documents.permissions import get_groups_with_only_permission
 from documents.permissions import set_permissions_for_object
 
 from . import bulk_edit
-from .models import DOCUMENT_SOURCE
-from .models import ConsumptionTemplate
 from .models import Correspondent
 from .models import Document
 from .models import DocumentType
@@ -1044,7 +1043,7 @@ class BulkEditObjectPermissionsSerializer(serializers.Serializer, SetPermissions
 class ConsumptionTemplateSerializer(OwnedObjectSerializer):
     order = serializers.IntegerField(required=False)
     sources = fields.MultipleChoiceField(
-        choices=DOCUMENT_SOURCE,
+        choices=ConsumptionTemplate.DocumentSourceChoices.choices,
         allow_empty=False,
         default={
             DocumentSource.ConsumeFolder,
@@ -1084,8 +1083,8 @@ class ConsumptionTemplateSerializer(OwnedObjectSerializer):
         ]
 
     def validate(self, attrs):
-        if ("filter_mailrule") in attrs:
-            attrs["sources"] = {int(DocumentSource.MailFetch)}
+        if ("filter_mailrule") in attrs and attrs["filter_mailrule"] is not None:
+            attrs["sources"] = {DocumentSource.MailFetch.value}
         if (
             ("filter_mailrule" not in attrs)
             and ("filter_filename" not in attrs or len(attrs["filter_filename"]) == 0)
diff --git a/src/documents/views.py b/src/documents/views.py
index e852738d7..4a9080b53 100644
--- a/src/documents/views.py
+++ b/src/documents/views.py
@@ -1253,10 +1253,12 @@ class BulkEditObjectPermissionsView(GenericAPIView, PassUserMixin):
 
 
 class ConsumptionTemplateViewSet(ModelViewSet, PassUserMixin):
+    permission_classes = (IsAuthenticated, PaperlessObjectPermissions)
+
+    serializer_class = ConsumptionTemplateSerializer
+    pagination_class = StandardPagination
+    filter_backends = (ObjectOwnedOrGrantedPermissionsFilter,)
+
     model = ConsumptionTemplate
 
     queryset = ConsumptionTemplate.objects.all().order_by("order")
-    serializer_class = ConsumptionTemplateSerializer
-    pagination_class = StandardPagination
-    permission_classes = (IsAuthenticated, PaperlessObjectPermissions)
-    filter_backends = (ObjectOwnedOrGrantedPermissionsFilter,)