Use generator in Mailbox.get_mail to limit memory leak (#149)

* Use generator in Mailbox.get_mail to limit memory leak

* Fix generator support in tests/base.py

django_mailbox/models.py

@@ -408,17 +408,16 @@ class Mailbox(models.Model):
         new_mail = []
         connection = self.get_connection()
         if not connection:
-            return new_mail
+            return
         for message in connection.get_message(condition):
             msg = self.process_incoming_message(message)
             if not msg is None:
-                new_mail.append(msg)
+                yield msg
         self.last_polling = now()
         if django.VERSION >= (1, 5):  # Django 1.5 introduces update_fields
             self.save(update_fields=['last_polling'])
         else:
             self.save()
-        return new_mail
 
     def __str__(self):
         return self.name

django_mailbox/tests/base.py

@@ -60,15 +60,25 @@ class EmailMessageTestCase(TestCase):
         super(EmailMessageTestCase, self).setUp()
 
     def _get_new_messages(self, mailbox, condition=None):
-        maximum_wait = time.time() + self.maximum_wait_seconds
+        start_time = time.time()
-        while True:
+        # wait until there is at least one message
-            if time.time() > maximum_wait:
+        while time.time() - start_time < self.maximum_wait_seconds:
-                raise EmailIntegrationTimeout()
+
             messages = self.mailbox.get_new_mail(condition)
-            if messages:
+
-                return messages
+            try:
+                # check if generator contains at least one element
+                message = next(messages)
+                yield message
+                for message in messages:
+                    yield message
+                return
+
+            except StopIteration:
                 time.sleep(5)
 
+        raise EmailIntegrationTimeout()
+
     def _get_email_as_text(self, name):
         with open(
             os.path.join(

django_mailbox/tests/test_integration_imap.py

@@ -64,7 +64,8 @@ class TestImap(EmailMessageTestCase):
             self.mailbox,
             condition=lambda m: m['subject'] == self.arbitrary_identifier
         )
+        message = next(messages)
 
-        self.assertEqual(1, len(messages))
+        self.assertEqual(message.subject, self.arbitrary_identifier)
-        self.assertEqual(messages[0].subject, self.arbitrary_identifier)
+        self.assertEqual(message.text, text_content)
-        self.assertEqual(messages[0].text, text_content)
+        self.assertEqual(0, len(list(messages)))

django_mailbox/tests/test_mailbox.py

@@ -32,5 +32,5 @@ class TestMailbox(TestCase):
                 'generic_message.eml',
             ))
         self.assertEqual(mailbox.last_polling, None)
-        mailbox.get_new_mail()
+        list(mailbox.get_new_mail())
         self.assertNotEqual(mailbox.last_polling, None)