Add more tests for MailMessageDecryptor

src/paperless_mail/tests/test_preprocessor.py

@@ -1,23 +1,20 @@
+import email
+import email.contentmanager
 import tempfile
+from email.message import Message
 from email.mime.application import MIMEApplication
 from email.mime.multipart import MIMEMultipart
-from typing import TYPE_CHECKING
 
 import gnupg
-from django.test import TestCase
 from django.test import override_settings
 from imap_tools import MailMessage
 
-from paperless_mail.mail import MailAccountHandler
 from paperless_mail.models import MailAccount
 from paperless_mail.models import MailRule
 from paperless_mail.preprocessor import MailMessageDecryptor
-from paperless_mail.tests.test_mail import MailMocker
+from paperless_mail.tests.test_mail import TestMail
 from paperless_mail.tests.test_mail import _AttachmentDef
 
-if TYPE_CHECKING:
-    import email.contentmanager
-
 
 class MessageEncryptor:
     def __init__(self):
@@ -35,10 +32,24 @@ class MessageEncryptor:
         )
         self.gpg.gen_key(input_data)
 
+    @staticmethod
+    def get_email_body_without_headers(email_message: Message) -> bytes:
+        """
+        Filters some relevant headers from an EmailMessage and returns just the body.
+        """
+        message_copy = email.message_from_bytes(email_message.as_bytes())
+
+        message_copy._headers = [
+            header
+            for header in message_copy._headers
+            if header[0].lower() not in ("from", "to", "subject")
+        ]
+        return message_copy.as_bytes()
+
     def encrypt(self, message):
         original_email: email.message.Message = message.obj
         encrypted_data = self.gpg.encrypt(
-            original_email.as_bytes(),
+            self.get_email_body_without_headers(original_email),
             self._testUser,
             armor=True,
         )
@@ -65,44 +76,66 @@ class MessageEncryptor:
         )
         new_email.attach(encrypted_part)
 
-        encrypted_message = MailMessage(
+        encrypted_message: MailMessage = MailMessage(
             [(f"UID {message.uid}".encode(), new_email.as_bytes())],
         )
         return encrypted_message
 
 
-class TestPreprocessor(TestCase):
+class TestMailMessageGpgDecryptor(TestMail):
     def setUp(self):
-        self.mailMocker = MailMocker()
-        self.mailMocker.setUp()
-
         self.messageEncryptor = MessageEncryptor()
         with override_settings(
             EMAIL_GNUPG_HOME=self.messageEncryptor.gpg_home,
             EMAIL_ENABLE_GPG_DECRYPTOR=True,
         ):
-            self.mail_account_handler = MailAccountHandler()
-
             super().setUp()
 
+    def test_preprocessor_is_able_to_run(self):
+        with override_settings(
+            EMAIL_GNUPG_HOME=self.messageEncryptor.gpg_home,
+            EMAIL_ENABLE_GPG_DECRYPTOR=True,
+        ):
+            self.assertTrue(MailMessageDecryptor.able_to_run())
+
+    def test_preprocessor_is_able_to_run2(self):
+        with override_settings(
+            EMAIL_GNUPG_HOME=None,
+            EMAIL_ENABLE_GPG_DECRYPTOR=True,
+        ):
+            self.assertTrue(MailMessageDecryptor.able_to_run())
+
+    def test_is_not_able_to_run_disabled(self):
+        with override_settings(
+            EMAIL_ENABLE_GPG_DECRYPTOR=False,
+        ):
+            self.assertFalse(MailMessageDecryptor.able_to_run())
+
+    def test_is_not_able_to_run_bogus_path(self):
+        with override_settings(
+            EMAIL_ENABLE_GPG_DECRYPTOR=True,
+            EMAIL_GNUPG_HOME="_)@# notapath &%#$",
+        ):
+            self.assertFalse(MailMessageDecryptor.able_to_run())
+
+    def test_decrypt_fails(self):
+        encrypted_message, _ = self.create_encrypted_unencrypted_message_pair()
+        empty_gpg_home = tempfile.mkdtemp()
+        with override_settings(
+            EMAIL_ENABLE_GPG_DECRYPTOR=True,
+            EMAIL_GNUPG_HOME=empty_gpg_home,
+        ):
+            message_decryptor = MailMessageDecryptor()
+            self.assertRaises(Exception, message_decryptor.run, encrypted_message)
+
     def test_decrypt_encrypted_mail(self):
         """
         Creates a mail with attachments. Then encrypts it with a new key.
         Verifies that this encrypted message can be decrypted with attachments intact.
         """
-        message = self.mailMocker.messageBuilder.create_message(
+        encrypted_message, message = self.create_encrypted_unencrypted_message_pair()
-            body="Test message with 2 attachments",
-            attachments=[
-                _AttachmentDef(
-                    filename="f1.pdf",
-                    disposition="inline",
-                ),
-                _AttachmentDef(filename="f2.pdf"),
-            ],
-        )
         headers = message.headers
         text = message.text
-        encrypted_message = self.messageEncryptor.encrypt(message)
 
         self.assertEqual(len(encrypted_message.attachments), 1)
         self.assertEqual(encrypted_message.attachments[0].filename, "encrypted.asc")
@@ -123,6 +156,20 @@ class TestPreprocessor(TestCase):
         self.assertEqual(decrypted_message.text, text)
         self.assertEqual(decrypted_message.uid, message.uid)
 
+    def create_encrypted_unencrypted_message_pair(self):
+        message = self.mailMocker.messageBuilder.create_message(
+            body="Test message with 2 attachments",
+            attachments=[
+                _AttachmentDef(
+                    filename="f1.pdf",
+                    disposition="inline",
+                ),
+                _AttachmentDef(filename="f2.pdf"),
+            ],
+        )
+        encrypted_message = self.messageEncryptor.encrypt(message)
+        return encrypted_message, message
+
     def test_handle_encrypted_message(self):
         message = self.mailMocker.messageBuilder.create_message(
             subject="the message title",