diff --git a/django_mailbox/admin.py b/django_mailbox/admin.py index 045877d..4d489b1 100644 --- a/django_mailbox/admin.py +++ b/django_mailbox/admin.py @@ -20,11 +20,7 @@ logger = logging.getLogger(__name__) def get_new_mail(mailbox_admin, request, queryset): - count = 0 - for mailbox in queryset.all(): - logger.debug('Receiving mail for %s' % mailbox) - count += sum(1 for i in mailbox.get_new_mail()) - logger.debug('Received %d %s.', count, 'mails' if count != 1 else 'mail') + queryset.get_new_mail() get_new_mail.short_description = _('Get new mail') diff --git a/django_mailbox/models.py b/django_mailbox/models.py index 4e5b111..69a1d0c 100644 --- a/django_mailbox/models.py +++ b/django_mailbox/models.py @@ -36,9 +36,23 @@ from django_mailbox.transports import Pop3Transport, ImapTransport, \ logger = logging.getLogger(__name__) +class MailboxQuerySet(models.QuerySet): + def get_new_mail(self): + count = 0 + for mailbox in self.all(): + logger.debug('Receiving mail for %s' % mailbox) + count += sum(1 for i in mailbox.get_new_mail()) + logger.debug('Received %d %s.', count, 'mails' if count != 1 else 'mail') + + +class MailboxManager(models.Manager): + def get_queryset(self): + return MailboxQuerySet(self.model, using=self._db) + + class ActiveMailboxManager(models.Manager): def get_queryset(self): - return super().get_queryset().filter( + return MailboxQuerySet(self.model, using=self._db).filter( active=True, ) @@ -106,7 +120,7 @@ class Mailbox(models.Model): null=True ) - objects = models.Manager() + objects = MailboxManager() active_mailboxes = ActiveMailboxManager() @property diff --git a/django_mailbox/tests/test_mailbox.py b/django_mailbox/tests/test_mailbox.py index b55e70c..e233f75 100644 --- a/django_mailbox/tests/test_mailbox.py +++ b/django_mailbox/tests/test_mailbox.py @@ -34,3 +34,13 @@ class TestMailbox(TestCase): self.assertEqual(mailbox.last_polling, None) list(mailbox.get_new_mail()) self.assertNotEqual(mailbox.last_polling, None) + + def test_queryset_get_new_mail(self): + mailbox = Mailbox.objects.create(uri="mbox://" + os.path.join( + os.path.dirname(__file__), + 'messages', + 'generic_message.eml', + )) + Mailbox.objects.filter(pk=mailbox.pk).get_new_mail() + mailbox.refresh_from_db() + self.assertNotEqual(mailbox.last_polling, None)