From ff919ae8f5bacd22ff1217d5c3f4f3036ff13a69 Mon Sep 17 00:00:00 2001
From: Mathieu Hinderyckx <mhindery@Mathieus-MacBook-Pro.local>
Date: Tue, 15 Nov 2016 21:12:43 +0100
Subject: [PATCH] Add management command + associated tests to cleanup old
 messages from the database.

---
 .../commands/djmail_delete_old_messages.py    | 20 ++++++++++
 djmail/tests.py                               | 37 +++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 djmail/management/commands/djmail_delete_old_messages.py

diff --git a/djmail/management/commands/djmail_delete_old_messages.py b/djmail/management/commands/djmail_delete_old_messages.py
new file mode 100644
index 0000000..3477ec1
--- /dev/null
+++ b/djmail/management/commands/djmail_delete_old_messages.py
@@ -0,0 +1,20 @@
+# -*- encoding: utf-8 -*-
+
+from datetime import datetime, timedelta
+from django.core.management.base import BaseCommand
+from djmail import models
+
+
+class Command(BaseCommand):
+    help = 'Remove (succesfully sent) messages older than specified amount of days.'
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            '--days',
+            type=int,
+            default=183, # default = 6 months
+            help='Number of days to use as cut-off for deletion',
+        )
+
+    def handle(self, *args, **options):
+        models.Message.objects.filter(sent_at__lt=datetime.now() - timedelta(days=options['days']), status=models.STATUS_SENT).delete()
diff --git a/djmail/tests.py b/djmail/tests.py
index 88c9b1f..3f1a818 100644
--- a/djmail/tests.py
+++ b/djmail/tests.py
@@ -3,11 +3,16 @@
 from __future__ import unicode_literals
 
 import json
+import sys
+from datetime import datetime, timedelta
+
 
 from django.core.mail import EmailMessage
+from django.core.management import call_command
 from django.core import mail
 from django.test import TestCase
 from django.test.utils import override_settings
+from django.utils.six import StringIO
 
 from . import core
 from . import models
@@ -280,3 +285,35 @@ def test_simple_send_email_with_magic_builder_1(self):
         self.assertEqual(email.to, model.to_email.split(','))
         self.assertEqual(email.subject, model.subject)
         self.assertEqual(email.body, model.body_text)
+
+
+class CleanupManagementCommand(EmailTestCaseMixin, TestCase):
+    def setUp(self):
+        super(CleanupManagementCommand, self).setUp()
+        # Create a message that was succesfully sent 1 year ago
+        self.old_log = models.Message.from_email_message(self.email)
+        self.old_log.status = models.STATUS_SENT
+        self.old_log.sent_at = datetime.now() - timedelta(days=365)
+        self.old_log.save()
+
+    @override_settings(
+        EMAIL_BACKEND='djmail.backends.default.EmailBackend',
+        DJMAIL_REAL_BACKEND='django.core.mail.backends.locmem.EmailBackend')
+    def test_delete_old_message_with_default_days(self):
+        self.email.send()
+        self.assertEqual(models.Message.objects.count(), 2)
+        out = StringIO()
+        sys.stout = out
+        call_command('djmail_delete_old_messages', stdout=out)
+        self.assertEqual(models.Message.objects.count(), 1)
+
+    @override_settings(
+        EMAIL_BACKEND='djmail.backends.default.EmailBackend',
+        DJMAIL_REAL_BACKEND='django.core.mail.backends.locmem.EmailBackend')
+    def test_retain_old_message_with_specified_days(self):
+        self.email.send()
+        self.assertEqual(models.Message.objects.count(), 2)
+        out = StringIO()
+        sys.stout = out
+        call_command('djmail_delete_old_messages', '--days', '366', stdout=out)
+        self.assertEqual(models.Message.objects.count(), 2)