From be2c55af384f2f3312482293e211efc6000d4b7b Mon Sep 17 00:00:00 2001
From: Trenton H <797416+stumpylog@users.noreply.github.com>
Date: Wed, 26 Jun 2024 11:04:07 -0700
Subject: [PATCH] Also converts the live mail account over (probably)

---
 src/paperless_mail/tests/conftest.py       | 24 ++++++++++
 src/paperless_mail/tests/test_live_mail.py | 53 ++++++++++------------
 2 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/src/paperless_mail/tests/conftest.py b/src/paperless_mail/tests/conftest.py
index dcfa904a1..01a98d57d 100644
--- a/src/paperless_mail/tests/conftest.py
+++ b/src/paperless_mail/tests/conftest.py
@@ -1,7 +1,11 @@
+import os
+from collections.abc import Generator
 from pathlib import Path
 
 import pytest
 
+from paperless_mail.mail import MailAccountHandler
+from paperless_mail.models import MailAccount
 from paperless_mail.parsers import MailDocumentParser
 
 
@@ -63,3 +67,23 @@ def merged_pdf_second(sample_dir: Path) -> Path:
 @pytest.fixture()
 def mail_parser() -> MailDocumentParser:
     return MailDocumentParser(logging_group=None)
+
+
+@pytest.fixture()
+def live_mail_account() -> Generator[MailAccount, None, None]:
+    try:
+        account = MailAccount.objects.create(
+            name="test",
+            imap_server=os.environ["PAPERLESS_MAIL_TEST_HOST"],
+            username=os.environ["PAPERLESS_MAIL_TEST_USER"],
+            password=os.environ["PAPERLESS_MAIL_TEST_PASSWD"],
+            imap_port=993,
+        )
+        yield account
+    finally:
+        account.delete()
+
+
+@pytest.fixture()
+def mail_account_handler() -> MailAccountHandler:
+    return MailAccountHandler()
diff --git a/src/paperless_mail/tests/test_live_mail.py b/src/paperless_mail/tests/test_live_mail.py
index 6de2a6770..ecf9f73b6 100644
--- a/src/paperless_mail/tests/test_live_mail.py
+++ b/src/paperless_mail/tests/test_live_mail.py
@@ -1,7 +1,7 @@
 import os
+import warnings
 
 import pytest
-from django.test import TestCase
 
 from paperless_mail.mail import MailAccountHandler
 from paperless_mail.mail import MailError
@@ -16,53 +16,46 @@ from paperless_mail.models import MailRule
     or not len(os.environ["PAPERLESS_MAIL_TEST_HOST"]),
     reason="Live server testing not enabled",
 )
-class TestMailLiveServer(TestCase):
-    def setUp(self) -> None:
-        self.mail_account_handler = MailAccountHandler()
-        self.account = MailAccount.objects.create(
-            name="test",
-            imap_server=os.environ["PAPERLESS_MAIL_TEST_HOST"],
-            username=os.environ["PAPERLESS_MAIL_TEST_USER"],
-            password=os.environ["PAPERLESS_MAIL_TEST_PASSWD"],
-            imap_port=993,
-        )
-
-        return super().setUp()
-
-    def tearDown(self) -> None:
-        self.account.delete()
-        return super().tearDown()
-
-    def test_process_non_gmail_server_flag(self):
+@pytest.mark.django_db()
+class TestMailLiveServer:
+    def test_process_non_gmail_server_flag(
+        self,
+        mail_account_handler: MailAccountHandler,
+        live_mail_account: MailAccount,
+    ):
         try:
             rule1 = MailRule.objects.create(
                 name="testrule",
-                account=self.account,
+                account=live_mail_account,
                 action=MailRule.MailAction.FLAG,
             )
 
-            self.mail_account_handler.handle_mail_account(self.account)
+            mail_account_handler.handle_mail_account(live_mail_account)
 
             rule1.delete()
 
         except MailError as e:
-            self.fail(f"Failure: {e}")
-        except Exception:
-            pass
+            pytest.fail(f"Failure: {e}")
+        except Exception as e:
+            warnings.warn(f"Unhandled exception: {e}")
 
-    def test_process_non_gmail_server_tag(self):
+    def test_process_non_gmail_server_tag(
+        self,
+        mail_account_handler: MailAccountHandler,
+        live_mail_account: MailAccount,
+    ):
         try:
             rule2 = MailRule.objects.create(
                 name="testrule",
-                account=self.account,
+                account=live_mail_account,
                 action=MailRule.MailAction.TAG,
             )
 
-            self.mail_account_handler.handle_mail_account(self.account)
+            mail_account_handler.handle_mail_account(live_mail_account)
 
             rule2.delete()
 
         except MailError as e:
-            self.fail(f"Failure: {e}")
-        except Exception:
-            pass
+            pytest.fail(f"Failure: {e}")
+        except Exception as e:
+            warnings.warn(f"Unhandled exception: {e}")