From e7e511cbe5b4056f3ffde1aab47c10e98260fa0b Mon Sep 17 00:00:00 2001 From: Adam Dobrawy Date: Wed, 23 Aug 2017 16:31:37 +0200 Subject: [PATCH] Fix generator support in tests/base.py --- django_mailbox/tests/base.py | 24 +++++++++++++------ django_mailbox/tests/test_integration_imap.py | 7 +++--- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/django_mailbox/tests/base.py b/django_mailbox/tests/base.py index 491f937..4884765 100644 --- a/django_mailbox/tests/base.py +++ b/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 - while True: - if time.time() > maximum_wait: - raise EmailIntegrationTimeout() + start_time = time.time() + # wait until there is at least one message + while time.time() - start_time < self.maximum_wait_seconds: + messages = self.mailbox.get_new_mail(condition) - if messages: - return list(messages) - time.sleep(5) + + 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( diff --git a/django_mailbox/tests/test_integration_imap.py b/django_mailbox/tests/test_integration_imap.py index 96c3a6e..1110f69 100644 --- a/django_mailbox/tests/test_integration_imap.py +++ b/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(messages[0].subject, self.arbitrary_identifier) - self.assertEqual(messages[0].text, text_content) + self.assertEqual(message.subject, self.arbitrary_identifier) + self.assertEqual(message.text, text_content) + self.assertEqual(0, len(list(messages)))