diff --git a/django_mailbox/admin.py b/django_mailbox/admin.py index e9a6ef7..4d489b1 100644 --- a/django_mailbox/admin.py +++ b/django_mailbox/admin.py @@ -20,9 +20,7 @@ logger = logging.getLogger(__name__) def get_new_mail(mailbox_admin, request, queryset): - for mailbox in queryset.all(): - logger.debug('Receiving mail for %s' % mailbox) - mailbox.get_new_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..5b7fbbc 100644 --- a/django_mailbox/models.py +++ b/django_mailbox/models.py @@ -36,7 +36,21 @@ from django_mailbox.transports import Pop3Transport, ImapTransport, \ logger = logging.getLogger(__name__) -class ActiveMailboxManager(models.Manager): +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(MailboxManager): def get_queryset(self): return super().get_queryset().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)