fix tests after rewrite

4 files changed, 227 insertions, 107 deletions

diff --git a/src/leap/mail/imap/fields.py b/src/leap/mail/imap/fields.py
index 40817cd..bc536fe 100644
--- a/src/leap/mail/imap/fields.py
+++ b/src/leap/mail/imap/fields.py
@@ -35,6 +35,7 @@ class WithMsgFields(object):
     UID_KEY = "uid"
     MBOX_KEY = "mbox"
     SEEN_KEY = "seen"
+    DEL_KEY = "deleted"
     RECENT_KEY = "recent"
     FLAGS_KEY = "flags"
     MULTIPART_KEY = "multi"
@@ -95,6 +96,7 @@ class WithMsgFields(object):
     TYPE_SUBS_IDX = 'by-type-and-subscribed'
     TYPE_MBOX_SEEN_IDX = 'by-type-and-mbox-and-seen'
     TYPE_MBOX_RECT_IDX = 'by-type-and-mbox-and-recent'
+    TYPE_MBOX_DEL_IDX = 'by-type-and-mbox-and-deleted'
     TYPE_C_HASH_IDX = 'by-type-and-contenthash'
     TYPE_C_HASH_PART_IDX = 'by-type-and-contenthash-and-partnumber'
     TYPE_P_HASH_IDX = 'by-type-and-payloadhash'
@@ -128,6 +130,7 @@ class WithMsgFields(object):
         # messages
         TYPE_MBOX_SEEN_IDX: [KTYPE, MBOX_VAL, 'bool(seen)'],
         TYPE_MBOX_RECT_IDX: [KTYPE, MBOX_VAL, 'bool(recent)'],
+        TYPE_MBOX_DEL_IDX: [KTYPE, MBOX_VAL, 'bool(deleted)'],
         TYPE_MBOX_RECT_SEEN_IDX: [KTYPE, MBOX_VAL,
                                   'bool(recent)', 'bool(seen)'],
     }
diff --git a/src/leap/mail/imap/mailbox.py b/src/leap/mail/imap/mailbox.py
index 5ea6f55..10087f6 100644
--- a/src/leap/mail/imap/mailbox.py
+++ b/src/leap/mail/imap/mailbox.py
@@ -390,18 +390,17 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         else:
             flags = tuple(str(flag) for flag in flags)
 
-        d = self._do_add_message(message, flags, date, uid_next)
+        d = self._do_add_message(message, flags=flags, date=date, uid=uid_next)
         d.addCallback(self._notify_new)
         return d
 
     @deferred
-    def _do_add_message(self, message, flags, date, uid_next):
+    def _do_add_message(self, message, flags, date, uid):
         """
         Calls to the messageCollection add_msg method (deferred to thread).
         Invoked from addMessage.
         """
-        self.messages.add_msg(message, flags=flags, date=date,
-                              uid=uid_next)
+        self.messages.add_msg(message, flags=flags, date=date, uid=uid)
 
     def _notify_new(self, *args):
         """
@@ -436,21 +435,29 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         # we should postpone the removal
         self._soledad.delete_doc(self._get_mbox())
 
-    @deferred
+    def _close_cb(self, result):
+        self.closed = True
+
+    def close(self):
+        """
+        Expunge and mark as closed
+        """
+        d = self.expunge()
+        d.addCallback(self._close_cb)
+        return d
+
+    def _expunge_cb(self, result):
+        return result
+
     def expunge(self):
         """
         Remove all messages flagged \\Deleted
         """
         if not self.isWriteable():
             raise imap4.ReadOnlyMailbox
-        deleted = []
-        for m in self.messages:
-            if self.DELETED_FLAG in m.getFlags():
-                self.messages.remove(m)
-                # XXX this would ve more efficient if we can just pass
-                # a sequence of uids.
-                deleted.append(m.getUID())
-        return deleted
+        d = self.messages.remove_all_deleted()
+        d.addCallback(self._expunge_cb)
+        return d
 
     @deferred
     def fetch(self, messages, uid):
@@ -603,14 +610,6 @@ class SoledadMailbox(WithMsgFields, MBoxParser):
         self._signal_unread_to_ui()
         return result
 
-    @deferred
-    def close(self):
-        """
-        Expunge and mark as closed
-        """
-        self.expunge()
-        self.closed = True
-
     # IMessageCopier
 
     @deferred
diff --git a/src/leap/mail/imap/messages.py b/src/leap/mail/imap/messages.py
index 47c40d5..80411f9 100644
--- a/src/leap/mail/imap/messages.py
+++ b/src/leap/mail/imap/messages.py
@@ -20,9 +20,11 @@ LeapMessage and MessageCollection.
 import copy
 import logging
 import StringIO
-from collections import namedtuple
+
+from collections import defaultdict, namedtuple
 
 from twisted.mail import imap4
+from twisted.internet import defer
 from twisted.python import log
 from u1db import errors as u1db_errors
 from zope.interface import implements
@@ -182,6 +184,7 @@ class MessageAttachment(object):
         if not self._msg:
             return {}
         headers = dict(self._msg.items())
+
         names = map(lambda s: s.upper(), names)
         if negate:
             cond = lambda key: key.upper() not in names
@@ -329,6 +332,7 @@ class LeapMessage(fields, MailParser, MBoxParser):
         doc.content[self.FLAGS_KEY] = flags
         doc.content[self.SEEN_KEY] = self.SEEN_FLAG in flags
         doc.content[self.RECENT_KEY] = self.RECENT_FLAG in flags
+        doc.content[self.DEL_KEY] = self.DELETED_FLAG in flags
         self._soledad.put_doc(doc)
 
     def addFlags(self, flags):
@@ -455,6 +459,7 @@ class LeapMessage(fields, MailParser, MBoxParser):
         headers = self._get_headers()
         if not headers:
             return {'content-type': ''}
+
         names = map(lambda s: s.upper(), names)
         if negate:
             cond = lambda key: key.upper() not in names
@@ -465,8 +470,8 @@ class LeapMessage(fields, MailParser, MBoxParser):
 
             # twisted imap server expects headers to be lowercase
         head = dict(
-            map(str, (key, value)) if key.lower() != "content-type"
-            else map(str, (key.lower(), value))
+            (str(key), map(str, value)) if key.lower() != "content-type"
+            else (str(key.lower(), map(str, value)))
             for (key, value) in head.items())
 
         # unpack and filter original dict by negate-condition
@@ -670,6 +675,9 @@ class LeapMessage(fields, MailParser, MBoxParser):
         # until we think about a good way of deorphaning.
         # Maybe a crawler of unreferenced docs.
 
+        uid = self._uid
+        print "removing...", uid
+
         fd = self._get_flags_doc()
         hd = self._get_headers_doc()
         #bd = self._get_body_doc()
@@ -682,7 +690,11 @@ class LeapMessage(fields, MailParser, MBoxParser):
             #docs.append(ad)
 
         for d in filter(None, docs):
-            self._soledad.delete_doc(d)
+            try:
+                self._soledad.delete_doc(d)
+            except Exception as exc:
+                logger.error(exc)
+        return uid
 
     def does_exist(self):
         """
@@ -849,6 +861,7 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
 
             fields.SEEN_KEY: False,
             fields.RECENT_KEY: True,
+            fields.DEL_KEY: False,
             fields.FLAGS_KEY: [],
             fields.MULTIPART_KEY: False,
             fields.SIZE_KEY: 0
@@ -921,7 +934,7 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
 
         self.soledad_writer = MessageProducer(
             SoledadDocWriter(soledad),
-            period=0.05)
+            period=0.02)
 
     def _get_empty_doc(self, _type=FLAGS_DOC):
         """
@@ -966,7 +979,9 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
             (self.FLAGS_DOC, self.HEADERS_DOC, self.BODY_DOC))
 
         msg = self._get_parsed_msg(raw)
-        headers = dict(msg)
+        headers = defaultdict(list)
+        for k, v in msg.items():
+            headers[k].append(v)
         raw_str = msg.as_string()
         chash = self._get_hash(msg)
         multi = msg.is_multipart()
@@ -987,7 +1002,8 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
                     inner_parts.append(p)
         else:
             body = msg.get_payload()
-        logger.debug("adding msg (multipart:%s)" % multi)
+        logger.debug("adding msg with uid %s (multipart:%s)" % (
+            uid, multi))
 
         # flags doc ---------------------------------------
         fd[self.MBOX_KEY] = self.mbox
@@ -998,26 +1014,33 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
         if flags:
             fd[self.FLAGS_KEY] = map(self._stringify, flags)
             fd[self.SEEN_KEY] = self.SEEN_FLAG in flags
-            fd[self.RECENT_KEY] = self.RECENT_FLAG in flags
+            fd[self.DEL_KEY] = self.DELETED_FLAG in flags
+            fd[self.RECENT_KEY] = True  # set always by default
 
         # headers doc ----------------------------------------
         hd[self.CONTENT_HASH_KEY] = chash
         hd[self.HEADERS_KEY] = headers
+
+        print "headers"
+        import pprint
+        pprint.pprint(headers)
+
         if not subject and self.SUBJECT_FIELD in headers:
-            hd[self.SUBJECT_KEY] = headers[self.SUBJECT_FIELD]
+            hd[self.SUBJECT_KEY] = first(headers[self.SUBJECT_FIELD])
         else:
             hd[self.SUBJECT_KEY] = subject
         if not date and self.DATE_FIELD in headers:
-            hd[self.DATE_KEY] = headers[self.DATE_FIELD]
+            hd[self.DATE_KEY] = first(headers[self.DATE_FIELD])
         else:
             hd[self.DATE_KEY] = date
         if multi:
+            # XXX fix for multipart nested case
             hd[self.NUM_PARTS_KEY] = len(msg.get_payload())
 
         # body doc
         bd[self.CONTENT_HASH_KEY] = chash
         bd[self.BODY_KEY] = body
-        # in an ideal world, we would not need to save a copy of the
+        # XXX in an ideal world, we would not need to save a copy of the
         # raw message. But we'll keep it until we can be sure that
         # we can rebuild the original message from the parts.
         bd[self.RAW_KEY] = raw_str
@@ -1062,14 +1085,29 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
             self.soledad_writer.put(ptuple(
                 mode=ptuple.ATTACHMENT_CREATE, payload=at))
 
-    def remove(self, msg):
+    def _remove_cb(self, result):
+        return result
+
+    def remove_all_deleted(self):
+        """
+        Removes all messages flagged as deleted.
         """
-        Removes a message.
+        delete_deferl = []
+        for msg in self.get_deleted():
+            delete_deferl.append(msg.remove())
+        d1 = defer.gatherResults(delete_deferl, consumeErrors=True)
+        d1.addCallback(self._remove_cb)
+        return d1
 
-        :param msg: a  Leapmessage instance
+    def remove(self, msg):
+        """
+        Remove a given msg.
+        :param msg: the message to be removed
         :type msg: LeapMessage
         """
-        msg.remove()
+        d = msg.remove()
+        d.addCallback(self._remove_cb)
+        return d
 
     # getters
 
@@ -1178,7 +1216,7 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
 
     def recent_iter(self):
         """
-        Get an iterator for the message docs with `recent` flag.
+        Get an iterator for the message UIDs with `recent` flag.
 
         :return: iterator through recent message docs
         :rtype: iterable
@@ -1210,6 +1248,30 @@ class MessageCollection(WithMsgFields, IndexedDB, MailParser, MBoxParser):
             fields.TYPE_FLAGS_VAL, self.mbox, '1')
         return count
 
+    # deleted messages
+
+    def deleted_iter(self):
+        """
+        Get an iterator for the message UIDs with `deleted` flag.
+
+        :return: iterator through deleted message docs
+        :rtype: iterable
+        """
+        return (doc.content[self.UID_KEY] for doc in
+                self._soledad.get_from_index(
+                    fields.TYPE_MBOX_DEL_IDX,
+                    fields.TYPE_FLAGS_VAL, self.mbox, '1'))
+
+    def get_deleted(self):
+        """
+        Get all messages with the `Deleted` flag.
+
+        :returns: a generator of LeapMessages
+        :rtype: generator
+        """
+        return (LeapMessage(self._soledad, docid, self.mbox)
+                for docid in self.deleted_iter())
+
     def __len__(self):
         """
         Returns the number of messages on this mailbox.
diff --git a/src/leap/mail/imap/tests/test_imap.py b/src/leap/mail/imap/tests/test_imap.py
index ea75854..e1bed8c 100644
--- a/src/leap/mail/imap/tests/test_imap.py
+++ b/src/leap/mail/imap/tests/test_imap.py
@@ -25,7 +25,7 @@ XXX add authors from the original twisted tests.
 @license: GPLv3, see included LICENSE file
 """
 # XXX review license of the original tests!!!
-from nose.twistedtools import deferred
+from email import parser
 
 try:
     from cStringIO import StringIO
@@ -36,9 +36,13 @@ import os
 import types
 import tempfile
 import shutil
+import time
+
+from itertools import chain
 
 
 from mock import Mock
+from nose.twistedtools import deferred, stop_reactor
 
 
 from twisted.mail import imap4
@@ -58,9 +62,9 @@ import twisted.cred.portal
 # import u1db
 
 from leap.common.testing.basetest import BaseLeapTest
-from leap.mail.imap.server import SoledadMailbox
-from leap.mail.imap.server import SoledadBackedAccount
-from leap.mail.imap.server import MessageCollection
+from leap.mail.imap.account import SoledadBackedAccount
+from leap.mail.imap.mailbox import SoledadMailbox
+from leap.mail.imap.messages import MessageCollection
 
 from leap.soledad.client import Soledad
 from leap.soledad.client import SoledadCrypto
@@ -321,6 +325,9 @@ class IMAP4HelperMixin(BaseLeapTest):
         for mb in self.server.theAccount.mailboxes:
             self.server.theAccount.delete(mb)
 
+        # email parser
+        self.parser = parser.Parser()
+
     def tearDown(self):
         """
         tearDown method called after each test.
@@ -389,6 +396,7 @@ class MessageCollectionTestCase(IMAP4HelperMixin, unittest.TestCase):
     """
     Tests for the MessageCollection class
     """
+    count = 0
 
     def setUp(self):
         """
@@ -396,34 +404,35 @@ class MessageCollectionTestCase(IMAP4HelperMixin, unittest.TestCase):
         We override mixin method since we are only testing
         MessageCollection interface in this particular TestCase
         """
-        self.messages = MessageCollection("testmbox", self._soledad)
-        for m in self.messages.get_all():
-            self.messages.remove(m)
+        self.messages = MessageCollection("testmbox%s" % (self.count,),
+                                          self._soledad)
+        MessageCollectionTestCase.count += 1
 
     def tearDown(self):
         """
         tearDown method for each test
-        Delete the message collection
         """
         del self.messages
 
+    def wait(self):
+        time.sleep(2)
+
     def testEmptyMessage(self):
         """
         Test empty message and collection
         """
-        em = self.messages._get_empty_msg()
+        em = self.messages._get_empty_doc()
         self.assertEqual(
             em,
             {
-                "date": '',
                 "flags": [],
-                "headers": {},
                 "mbox": "inbox",
-                "raw": "",
                 "recent": True,
                 "seen": False,
-                "subject": "",
-                "type": "msg",
+                "deleted": False,
+                "multi": False,
+                "size": 0,
+                "type": "flags",
                 "uid": 1,
             })
         self.assertEqual(self.messages.count(), 0)
@@ -432,23 +441,22 @@ class MessageCollectionTestCase(IMAP4HelperMixin, unittest.TestCase):
         """
         Add multiple messages
         """
+        # TODO really profile addition
         mc = self.messages
+        print "messages", self.messages
         self.assertEqual(self.messages.count(), 0)
-        mc.add_msg('Stuff', subject="test1")
-        self.assertEqual(self.messages.count(), 1)
-        mc.add_msg('Stuff', subject="test2")
-        self.assertEqual(self.messages.count(), 2)
-        mc.add_msg('Stuff', subject="test3")
-        self.assertEqual(self.messages.count(), 3)
-        mc.add_msg('Stuff', subject="test4")
+        mc.add_msg('Stuff', uid=1, subject="test1")
+        mc.add_msg('Stuff', uid=2, subject="test2")
+        mc.add_msg('Stuff', uid=3, subject="test3")
+        mc.add_msg('Stuff', uid=4, subject="test4")
+        self.wait()
         self.assertEqual(self.messages.count(), 4)
-        mc.add_msg('Stuff', subject="test5")
-        mc.add_msg('Stuff', subject="test6")
-        mc.add_msg('Stuff', subject="test7")
-        mc.add_msg('Stuff', subject="test8")
-        mc.add_msg('Stuff', subject="test9")
-        mc.add_msg('Stuff', subject="test10")
-        self.assertEqual(self.messages.count(), 10)
+        mc.add_msg('Stuff', uid=5, subject="test5")
+        mc.add_msg('Stuff', uid=6, subject="test6")
+        mc.add_msg('Stuff', uid=7, subject="test7")
+        self.wait()
+        self.assertEqual(self.messages.count(), 7)
+        self.wait()
 
     def testRecentCount(self):
         """
@@ -456,45 +464,48 @@ class MessageCollectionTestCase(IMAP4HelperMixin, unittest.TestCase):
         """
         mc = self.messages
         self.assertEqual(self.messages.count_recent(), 0)
-        mc.add_msg('Stuff', subject="test1", uid=1)
+        mc.add_msg('Stuff', uid=1, subject="test1")
         # For the semantics defined in the RFC, we auto-add the
         # recent flag by default.
+        self.wait()
         self.assertEqual(self.messages.count_recent(), 1)
-        mc.add_msg('Stuff', subject="test2", uid=2, flags=('\\Deleted',))
+        mc.add_msg('Stuff', subject="test2", uid=2,
+                   flags=('\\Deleted',))
+        self.wait()
         self.assertEqual(self.messages.count_recent(), 2)
-        mc.add_msg('Stuff', subject="test3", uid=3, flags=('\\Recent',))
+        mc.add_msg('Stuff', subject="test3", uid=3,
+                   flags=('\\Recent',))
+        self.wait()
         self.assertEqual(self.messages.count_recent(), 3)
         mc.add_msg('Stuff', subject="test4", uid=4,
                    flags=('\\Deleted', '\\Recent'))
+        self.wait()
         self.assertEqual(self.messages.count_recent(), 4)
 
-        for m in mc:
-            msg = self.messages.get_msg_by_uid(m.get('uid'))
-            msg_newflags = msg.removeFlags(('\\Recent',))
-            self._soledad.put_doc(msg_newflags)
-
+        for msg in mc:
+            msg.removeFlags(('\\Recent',))
         self.assertEqual(mc.count_recent(), 0)
 
     def testFilterByMailbox(self):
         """
         Test that queries filter by selected mailbox
         """
+        def wait():
+            time.sleep(1)
+
         mc = self.messages
         self.assertEqual(self.messages.count(), 0)
-        mc.add_msg('', subject="test1")
-        self.assertEqual(self.messages.count(), 1)
-        mc.add_msg('', subject="test2")
-        self.assertEqual(self.messages.count(), 2)
-        mc.add_msg('', subject="test3")
+        mc.add_msg('', uid=1, subject="test1")
+        mc.add_msg('', uid=2, subject="test2")
+        mc.add_msg('', uid=3, subject="test3")
+        wait()
         self.assertEqual(self.messages.count(), 3)
-
-        newmsg = mc._get_empty_msg()
+        newmsg = mc._get_empty_doc()
         newmsg['mailbox'] = "mailbox/foo"
-        newmsg['subject'] = "test another mailbox"
         mc._soledad.create_doc(newmsg)
         self.assertEqual(mc.count(), 3)
         self.assertEqual(
-            len(mc._soledad.get_from_index(mc.TYPE_IDX, "*")), 4)
+            len(mc._soledad.get_from_index(mc.TYPE_IDX, "flags")), 4)
 
 
 class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
@@ -1174,16 +1185,20 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
         def login():
             return self.client.login('testuser', 'password-test')
 
+        def wait():
+            time.sleep(0.5)
+
         def append():
             return self.client.append(
                 'root/subthing',
                 message,
-                ['\\SEEN', '\\DELETED'],
+                ('\\SEEN', '\\DELETED'),
                 'Tue, 17 Jun 2003 11:22:16 -0600 (MDT)',
             )
 
         d1 = self.connected.addCallback(strip(login))
         d1.addCallbacks(strip(append), self._ebGeneral)
+        d1.addCallbacks(strip(wait), self._ebGeneral)
         d1.addCallbacks(self._cbStopClient, self._ebGeneral)
         d2 = self.loopback()
         d = defer.gatherResults([d1, d2])
@@ -1191,17 +1206,31 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
 
     def _cbTestFullAppend(self, ignored, infile):
         mb = SimpleLEAPServer.theAccount.getMailbox('root/subthing')
+        time.sleep(0.5)
         self.assertEqual(1, len(mb.messages))
 
+        msg = mb.messages.get_msg_by_uid(1)
         self.assertEqual(
-            ['\\SEEN', '\\DELETED'],
-            mb.messages[1].content['flags'])
+            ('\\SEEN', '\\DELETED'),
+            msg.getFlags())
 
         self.assertEqual(
             'Tue, 17 Jun 2003 11:22:16 -0600 (MDT)',
-            mb.messages[1].content['date'])
+            msg.getInternalDate())
+
+        parsed = self.parser.parse(open(infile))
+        body = parsed.get_payload()
+        headers = parsed.items()
+        self.assertEqual(
+            body,
+            msg.getBodyFile().read())
+
+        msg_headers = msg.getHeaders(True, "",)
+        gotheaders = list(chain(
+            *[[(k, item) for item in v] for (k, v) in msg_headers.items()]))
 
-        self.assertEqual(open(infile).read(), mb.messages[1].content['raw'])
+        self.assertItemsEqual(
+            headers, gotheaders)
 
     @deferred(timeout=None)
     def testPartialAppend(self):
@@ -1209,12 +1238,14 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
         Test partially appending a message to the mailbox
         """
         infile = util.sibpath(__file__, 'rfc822.message')
-        message = open(infile)
         SimpleLEAPServer.theAccount.addMailbox('PARTIAL/SUBTHING')
 
         def login():
             return self.client.login('testuser', 'password-test')
 
+        def wait():
+            time.sleep(1)
+
         def append():
             message = file(infile)
             return self.client.sendCommand(
@@ -1226,6 +1257,7 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
                 )
             )
         d1 = self.connected.addCallback(strip(login))
+        d1.addCallbacks(strip(wait), self._ebGeneral)
         d1.addCallbacks(strip(append), self._ebGeneral)
         d1.addCallbacks(self._cbStopClient, self._ebGeneral)
         d2 = self.loopback()
@@ -1235,15 +1267,20 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
 
     def _cbTestPartialAppend(self, ignored, infile):
         mb = SimpleLEAPServer.theAccount.getMailbox('PARTIAL/SUBTHING')
-
+        time.sleep(1)
         self.assertEqual(1, len(mb.messages))
+        msg = mb.messages.get_msg_by_uid(1)
         self.assertEqual(
-            ['\\SEEN', ],
-            mb.messages[1].content['flags']
+            ('\\SEEN', ),
+            msg.getFlags()
         )
+        #self.assertEqual(
+            #'Right now', msg.getInternalDate())
+        parsed = self.parser.parse(open(infile))
+        body = parsed.get_payload()
         self.assertEqual(
-            'Right now', mb.messages[1].content['date'])
-        self.assertEqual(open(infile).read(), mb.messages[1].content['raw'])
+            body,
+            msg.getBodyFile().read())
 
     @deferred(timeout=None)
     def testCheck(self):
@@ -1279,14 +1316,19 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
         self.server.theAccount.addMailbox(name)
 
         m = SimpleLEAPServer.theAccount.getMailbox(name)
-        m.messages.add_msg('', subject="Message 1",
+        m.messages.add_msg('test 1', uid=1, subject="Message 1",
                            flags=('\\Deleted', 'AnotherFlag'))
-        m.messages.add_msg('', subject="Message 2", flags=('AnotherFlag',))
-        m.messages.add_msg('', subject="Message 3", flags=('\\Deleted',))
+        m.messages.add_msg('test 2', uid=2, subject="Message 2",
+                           flags=('AnotherFlag',))
+        m.messages.add_msg('test 3', uid=3, subject="Message 3",
+                           flags=('\\Deleted',))
 
         def login():
             return self.client.login('testuser', 'password-test')
 
+        def wait():
+            time.sleep(1)
+
         def select():
             return self.client.select(name)
 
@@ -1294,6 +1336,7 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
             return self.client.close()
 
         d = self.connected.addCallback(strip(login))
+        d.addCallbacks(strip(wait), self._ebGeneral)
         d.addCallbacks(strip(select), self._ebGeneral)
         d.addCallbacks(strip(close), self._ebGeneral)
         d.addCallbacks(self._cbStopClient, self._ebGeneral)
@@ -1302,8 +1345,10 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
 
     def _cbTestClose(self, ignored, m):
         self.assertEqual(len(m.messages), 1)
+        messages = [msg for msg in m.messages]
+        self.assertFalse(messages[0] is None)
         self.assertEqual(
-            m.messages[1].content['subject'],
+            messages[0]._hdoc.content['subject'],
             'Message 2')
         self.failUnless(m.closed)
 
@@ -1315,17 +1360,19 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
         name = 'mailbox-expunge'
         SimpleLEAPServer.theAccount.addMailbox(name)
         m = SimpleLEAPServer.theAccount.getMailbox(name)
-        m.messages.add_msg('', subject="Message 1",
+        m.messages.add_msg('test 1', uid=1, subject="Message 1",
                            flags=('\\Deleted', 'AnotherFlag'))
-        self.failUnless(m.messages.count() == 1)
-        m.messages.add_msg('', subject="Message 2", flags=('AnotherFlag',))
-        self.failUnless(m.messages.count() == 2)
-        m.messages.add_msg('', subject="Message 3", flags=('\\Deleted',))
-        self.failUnless(m.messages.count() == 3)
+        m.messages.add_msg('test 2', uid=2, subject="Message 2",
+                           flags=('AnotherFlag',))
+        m.messages.add_msg('test 3', uid=3, subject="Message 3",
+                           flags=('\\Deleted',))
 
         def login():
             return self.client.login('testuser', 'password-test')
 
+        def wait():
+            time.sleep(2)
+
         def select():
             return self.client.select('mailbox-expunge')
 
@@ -1338,6 +1385,7 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
 
         self.results = None
         d1 = self.connected.addCallback(strip(login))
+        d1.addCallbacks(strip(wait), self._ebGeneral)
         d1.addCallbacks(strip(select), self._ebGeneral)
         d1.addCallbacks(strip(expunge), self._ebGeneral)
         d1.addCallbacks(expunged, self._ebGeneral)
@@ -1348,12 +1396,13 @@ class LeapIMAP4ServerTestCase(IMAP4HelperMixin, unittest.TestCase):
 
     def _cbTestExpunge(self, ignored, m):
         # we only left 1 mssage with no deleted flag
-        self.assertEqual(m.messages.count(), 1)
+        self.assertEqual(len(m.messages), 1)
+        messages = [msg for msg in m.messages]
         self.assertEqual(
-            m.messages[1].content['subject'],
+            messages[0]._hdoc.content['subject'],
             'Message 2')
-        self.assertEqual(self.results, [0, 1])
-        # XXX fix this thing with the indexes...
+        # the uids of the deleted messages
+        self.assertItemsEqual(self.results, [1, 3])
 
 
 class IMAP4ServerSearchTestCase(IMAP4HelperMixin, unittest.TestCase):
@@ -1363,3 +1412,10 @@ class IMAP4ServerSearchTestCase(IMAP4HelperMixin, unittest.TestCase):
     """
     # XXX coming soon to your screens!
     pass
+
+
+def tearDownModule():
+    """
+    Tear down functions for module level
+    """
+    stop_reactor()