Fix generator support in tests/base.py

django_mailbox/tests/base.py

@@ -60,14 +60,24 @@ 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 list(messages)
+            try:
-            time.sleep(5)
+                # 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(

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)))