make direct_imap idle usage safe from tests, so we never miss finishing direct_imap idle

python/src/deltachat/direct_imap.py

@@ -6,6 +6,7 @@ and for cleaning up inbox/mvbox for each test function run.
 import io
 import ssl
 import pathlib
+from contextlib import contextmanager
 from imap_tools import MailBox, MailBoxTls, errors, AND, Header, MailMessageFlags, MailMessage
 import imaplib
 from deltachat import const, Account
@@ -52,10 +53,6 @@ class DirectImap:
         self.select_folder("INBOX")
 
     def shutdown(self):
-        try:
-            self.idle_done()
-        except (OSError, imaplib.IMAP4.abort):
-            pass
         try:
             self.conn.logout()
         except (OSError, imaplib.IMAP4.abort):
@@ -151,49 +148,14 @@ class DirectImap:
 
         print(stream.getvalue(), file=logfile)
 
-    def idle_start(self):
+    @contextmanager
-        """ switch this connection to idle mode. non-blocking. """
+    def idle(self):
-        assert not self._idling
+        """ return Idle ContextManager. """
-        res = self.conn.idle.start()
+        idle_manager = IdleManager(self)
-        self._idling = True
+        try:
-        return res
+            yield idle_manager
-
+        finally:
-    def idle_check(self, terminate=False, timeout=None) -> List[bytes]:
+            idle_manager.done()
-        """ (blocking) wait for next idle message from server. """
-        assert self._idling
-        self.account.log("imap-direct: calling idle_check")
-        res = self.conn.idle.poll(timeout=timeout)
-        if terminate:
-            self.idle_done()
-        self.account.log("imap-direct: idle_check returned {!r}".format(res))
-        return res
-
-    def idle_wait_for_new_message(self, terminate=False, timeout=None) -> bytes:
-        while 1:
-            for item in self.idle_check(timeout=timeout):
-                if b'EXISTS' in item or b'RECENT' in item:
-                    if terminate:
-                        self.idle_done()
-                    return item
-
-    def idle_wait_for_seen(self, terminate=False, timeout=None) -> int:
-        """ Return first message with SEEN flag from a running idle-stream.
-        """
-        while 1:
-            for item in self.idle_check(timeout=timeout):
-                if FETCH in item:
-                    self.account.log(str(item))
-                    if FLAGS in item and rb'\Seen' in item:
-                        if terminate:
-                            self.idle_done()
-                        return int(item.split(b' ')[1])
-
-    def idle_done(self):
-        """ send idle-done to server if we are currently in idle mode. """
-        if self._idling:
-            res = self.conn.idle.stop()
-            self._idling = False
-            return res
 
     def append(self, folder: str, msg: str):
         """Upload a message to *folder*.
@@ -209,3 +171,38 @@ class DirectImap:
         if len(msgs) == 0:
             raise Exception("Did not find message " + message_id + ", maybe you forgot to select the correct folder?")
         return msgs[0]
+
+
+class IdleManager:
+    def __init__(self, direct_imap):
+        self.direct_imap = direct_imap
+        self.log = direct_imap.account.log
+        self.direct_imap.conn.idle.start()
+
+    def check(self, timeout=None) -> List[bytes]:
+        """ (blocking) wait for next idle message from server. """
+        self.log("imap-direct: calling idle_check")
+        res = self.direct_imap.conn.idle.poll(timeout=timeout)
+        self.log("imap-direct: idle_check returned {!r}".format(res))
+        return res
+
+    def wait_for_new_message(self, timeout=None) -> bytes:
+        while 1:
+            for item in self.check(timeout=timeout):
+                if b'EXISTS' in item or b'RECENT' in item:
+                    return item
+
+    def wait_for_seen(self, timeout=None) -> int:
+        """ Return first message with SEEN flag from a running idle-stream.
+        """
+        while 1:
+            for item in self.check(timeout=timeout):
+                if FETCH in item:
+                    self.log(str(item))
+                    if FLAGS in item and rb'\Seen' in item:
+                        return int(item.split(b' ')[1])
+
+    def done(self):
+        """ send idle-done to server if we are currently in idle mode. """
+        res = self.direct_imap.conn.idle.stop()
+        return res

python/tests/test_account.py

@@ -644,11 +644,11 @@ def test_basic_imap_api(acfactory, tmpdir):
 
     imap2 = ac2.direct_imap
 
-    imap2.idle_start()
+    with imap2.idle() as idle2:
         chat12.send_text("hello")
         ac2._evtracker.wait_next_incoming_message()
+        idle2.wait_for_new_message()
 
-    imap2.idle_check(terminate=True)
     assert imap2.get_unread_cnt() == 1
     imap2.mark_all_read()
     assert imap2.get_unread_cnt() == 0
@@ -745,7 +745,7 @@ class TestOnlineAccount:
         ac1.set_config("bcc_self", "1")
 
         lp.sec("send out message with bcc to ourselves")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
             msg_out = chat.send_text("message2")
 
             # wait for send out (BCC)
@@ -755,7 +755,7 @@ class TestOnlineAccount:
             # now make sure we are sending message to ourselves too
             assert self_addr in ev.data2
             assert other_addr in ev.data2
-        assert ac1.direct_imap.idle_wait_for_seen()
+            assert idle1.wait_for_seen()
 
         # Second client receives only second message, but not the first
         ev_msg = ac1_clone._evtracker.wait_next_messages_changed()
@@ -1026,23 +1026,22 @@ class TestOnlineAccount:
         msg4 = ac2._evtracker.wait_next_incoming_message()
 
         lp.sec("mark messages as seen on ac2, wait for changes on ac1")
-        ac2.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
-        ac1.direct_imap.idle_start()
+            with ac2.direct_imap.idle() as idle2:
-
                 ac2.mark_seen_messages([msg2, msg4])
                 ev = ac2._evtracker.get_matching("DC_EVENT_MSGS_NOTICED")
                 assert msg2.chat.id == msg4.chat.id
                 assert ev.data1 == msg2.chat.id
                 assert ev.data2 == 0
 
-        ac2.direct_imap.idle_wait_for_new_message(terminate=True)
+                idle2.wait_for_new_message()
             lp.step("1")
             for i in range(2):
                 ev = ac1._evtracker.get_matching("DC_EVENT_MSG_READ")
                 assert ev.data1 > const.DC_CHAT_ID_LAST_SPECIAL
                 assert ev.data2 > const.DC_MSG_ID_LAST_SPECIAL
             lp.step("2")
-        ac1.direct_imap.idle_wait_for_seen()  # Check that ac1 marks the read receipt as read
+            idle1.wait_for_seen()  # Check that ac1 marks the read receipt as read
 
         assert msg1.is_out_mdn_received()
         assert msg3.is_out_mdn_received()
@@ -1062,24 +1061,22 @@ class TestOnlineAccount:
         acfactory.bring_accounts_online()
 
         ac2.stop_io()
-        ac2.direct_imap.idle_start()
+        with ac2.direct_imap.idle() as idle2:
-
             ac1.create_chat(ac2).send_text("Hello!")
-
+            idle2.wait_for_new_message()
-        ac2.direct_imap.idle_wait_for_new_message(terminate=True)
 
         # Emulate moving of the message to DeltaChat folder by Sieve rule.
         ac2.direct_imap.conn.move(["*"], "DeltaChat")
-
         ac2.direct_imap.select_folder("DeltaChat")
-        ac2.direct_imap.idle_start()
+
+        with ac2.direct_imap.idle() as idle2:
             ac2.start_io()
             msg = ac2._evtracker.wait_next_incoming_message()
 
             # Accept the contact request.
             msg.chat.accept()
             ac2.mark_seen_messages([msg])
-        uid = ac2.direct_imap.idle_wait_for_seen(terminate=True)
+            uid = idle2.wait_for_seen()
 
         assert len([a for a in ac2.direct_imap.conn.fetch(AND(seen=True, uid=U(uid, "*")))]) == 1
 
@@ -1185,18 +1182,15 @@ class TestOnlineAccount:
         folder = "mvbox" if mvbox_move else "inbox"
         ac1.direct_imap.select_config_folder(folder)
         ac2.direct_imap.select_config_folder(folder)
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
-        ac2.direct_imap.idle_start()
+            with ac2.direct_imap.idle() as idle2:
-
                 acfactory.get_accepted_chat(ac1, ac2).send_text("hi")
                 msg = ac2._evtracker.wait_next_incoming_message()
 
                 ac2.mark_seen_messages([msg])
 
-        ac1.direct_imap.idle_wait_for_seen()  # Check that the mdn is marked as seen
+                idle2.wait_for_seen()  # Check original message is marked as seen
-        ac2.direct_imap.idle_wait_for_seen()  # Check that the original message is marked as seen
+                idle1.wait_for_seen()  # Check that the mdn is marked as seen
-        ac1.direct_imap.idle_done()
-        ac2.direct_imap.idle_done()
 
     def test_reply_privately(self, acfactory):
         ac1, ac2 = acfactory.get_online_accounts(2)
@@ -1249,8 +1243,7 @@ class TestOnlineAccount:
         assert len(msg.chat.get_messages()) == 1
 
         ac1.direct_imap.select_config_folder("mvbox")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
-
             lp.sec("ac2: mark incoming message as seen")
             ac2.mark_seen_messages([msg])
 
@@ -1261,7 +1254,7 @@ class TestOnlineAccount:
             assert len(chat.get_messages()) == 1
 
             # Wait for the message to be marked as seen on IMAP.
-        assert ac1.direct_imap.idle_wait_for_seen()
+            assert idle1.wait_for_seen()
 
         # MDN is received even though MDNs are already disabled
         assert msg_out.is_out_mdn_received()
@@ -2238,12 +2231,10 @@ class TestOnlineAccount:
                "all messages are fetched")
 
         ac1.direct_imap.select_config_folder("inbox")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
             ac2.create_chat(ac1).send_text("Hi")
-
+            idle1.wait_for_new_message()
-        ac1.direct_imap.idle_wait_for_new_message(terminate=True)
         ac1.maybe_network()
-
         ac1._evtracker.wait_for_all_work_done()
         msgs = ac1.create_chat(ac2).get_messages()
         assert len(msgs) == 1
@@ -2273,10 +2264,9 @@ class TestOnlineAccount:
         ac1.create_contact(ac2).block()
 
         ac1.direct_imap.select_config_folder("inbox")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
             ac2.create_chat(ac1).send_text("Hi")
-
+            idle1.wait_for_new_message()
-        ac1.direct_imap.idle_wait_for_new_message(terminate=True)
         ac1.maybe_network()
 
         while 1:
@@ -2743,9 +2733,9 @@ class TestOnlineAccount:
 
         lp.sec("Send a message to from ac2 to ac1 and manually move it to the mvbox")
         ac1.direct_imap.select_config_folder("inbox")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
             acfactory.get_accepted_chat(ac2, ac1).send_text("hello")
-        ac1.direct_imap.idle_wait_for_new_message(terminate=True)
+            idle1.wait_for_new_message()
         ac1.direct_imap.conn.move(["*"], folder)  # "*" means "biggest UID in mailbox"
 
         lp.sec("start_io() and see if DeltaChat finds the message (" + variant + ")")
@@ -2793,8 +2783,7 @@ class TestOnlineAccount:
         assert_folders_configured(ac1)
 
         assert ac1.direct_imap.select_config_folder("mvbox" if mvbox_move else "inbox")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
-
             lp.sec("send out message with bcc to ourselves")
             ac1.set_config("bcc_self", "1")
             chat = acfactory.get_accepted_chat(ac1, ac2)
@@ -2802,7 +2791,7 @@ class TestOnlineAccount:
             assert_folders_configured(ac1)
 
             lp.sec("wait until the bcc_self message arrives in correct folder and is marked seen")
-        assert ac1.direct_imap.idle_wait_for_seen()
+            assert idle1.wait_for_seen()
         assert_folders_configured(ac1)
 
         lp.sec("create a cloned ac1 and fetch contact history during configure")
@@ -2843,13 +2832,12 @@ class TestOnlineAccount:
         ac1._evtracker.wait_next_incoming_message()
 
         lp.sec("send out message with bcc to ourselves")
-        ac1.direct_imap.idle_start()
+        with ac1.direct_imap.idle() as idle1:
             ac1.set_config("bcc_self", "1")
             ac1_ac2_chat = ac1.create_chat(ac2)
             ac1_ac2_chat.send_text("outgoing, encrypted direct message, creating a chat")
-
+            # wait until the bcc_self message arrives
-        # now wait until the bcc_self message arrives
+            assert idle1.wait_for_seen()
-        assert ac1.direct_imap.idle_wait_for_seen()
 
         lp.sec("Clone online account and let it fetch the existing messages")
         ac1_clone = acfactory.new_online_configuring_account(cloned_from=ac1)