IMAP source improvements

- Support multiple STORE commands. Defaults to +FLAGS (\Deleted) to conserve current behavior. - Support custom search criteria. Defaults to NOT DELETED to converse current behavior. - Support enabling/disabling IMAP expunge command after each message. Defaults to True to conserve current behavior. - Increase functionality of the MockIMAPServer Signed-off-by: Roberto Rosario <roberto.rosario@mayan-edms.com>
2019-10-24 04:02:25 -04:00
parent 1b6468522a
commit 4dea4129db
10 changed files with 294 additions and 54 deletions
@@ -260,44 +260,136 @@ class EmailBaseTestCase(GenericDocumentTestCase):
            )


+class MockIMAPMessage(object):
+    flags = []
+    uid = None
+
+    def __init__(self):
+        self.uid = '999'
+
+    def get_flags(self):
+        return ' '.join(self.flags)
+
+
+class MockIMAPMailbox(object):
+    messages = []
+
+    def __init__(self, name='INBOX'):
+        self.name = name
+
+
+class MockIMAPServer(object):
+    def __init__(self):
+        self.mailboxes = {
+            'INBOX': MockIMAPMailbox(name='INBOX')
+        }
+        self.mailboxes['INBOX'].messages.append(MockIMAPMessage())
+        self.mailbox_selected = None
+
+    def login(self, user, password):
+        return ('OK', ['{} authenticated (Success)'.format(user)])
+
+    def select(self, mailbox='INBOX', readonly=False):
+        self.mailbox_selected = self.mailboxes[mailbox]
+
+        return (
+            'OK', [
+                len(self.mailbox_selected.messages)
+            ]
+        )
+
+    def search(self, charset, *criteria):
+        """
+        7.2.5.  SEARCH Response
+        Contents:   zero or more numbers
+        The SEARCH response occurs as a result of a SEARCH or UID SEARCH
+        command.  The number(s) refer to those messages that match the
+        search criteria.  For SEARCH, these are message sequence numbers;
+        for UID SEARCH, these are unique identifiers.  Each number is
+        delimited by a space.
+
+        Example:    S: * SEARCH 2 3 6
+        """
+        results = [
+            self.mailbox_selected.messages[0]
+        ]
+
+        message_sequences = []
+        for result in results:
+            message_sequences.append(
+                force_text(
+                    self.mailbox_selected.messages.index(result)
+                )
+            )
+
+        return ('OK', ' '.join(message_sequences))
+
+    def fetch(self, message_set, message_parts):
+        results = []
+        for message_number in message_set.split():
+            message = self.mailbox_selected.messages[int(message_number) - 1]
+            if '\\Seen' not in message.flags:
+                message.flags.append('\\Seen')
+
+            results.append(
+                (
+                    '{} (RFC822 {{4800}}'.format(message_number),
+                    TEST_EMAIL_BASE64_FILENAME,
+                    ' FLAGS ({}))'.format(message.get_flags())
+                )
+            )
+        return ('OK', results)
+
+    def store(self, message_set, command, flags):
+        results = []
+
+        for message_number in message_set.split():
+            message = self.mailbox_selected.messages[int(message_number) - 1]
+
+            if command == 'FLAGS':
+                message.flags = flags.split()
+            elif command == '+FLAGS':
+                for flag in flags.split():
+                    if flag not in message.flags:
+                        message.flags.append(flag)
+            elif command == '-FLAGS':
+                for flag in flags.split():
+                    if flag in message.flags:
+                        message.flags.remove(flag)
+
+            results.append(
+                '{} (FLAGS ({}))'.format(message_number, message.get_flags())
+            )
+
+        return ('OK', results)
+
+    def expunge(self):
+        result = []
+
+        for message in self.mailbox_selected.messages:
+            if '\\Deleted' in message.flags:
+                result.append(
+                    force_text(
+                        self.mailbox_selected.messages.index(message)
+                    )
+                )
+                self.mailbox_selected.messages.remove(message)
+
+        return ('OK', ' '.join(result))
+
+    def close(self):
+        return ('OK', ['Returned to authenticated state. (Success)'])
+
+    def logout(self):
+        return ('BYE', ['LOGOUT Requested'])
+
+
 class IMAPSourceTestCase(GenericDocumentTestCase):
    auto_upload_document = False

-    class MockIMAPServer(object):
-        def login(self, user, password):
-            return ('OK', ['{} authenticated (Success)'.format(user)])
-
-        def select(self, mailbox='INBOX', readonly=False):
-            return ('OK', ['1'])
-
-        def search(self, charset, *criteria):
-            return ('OK', ['1'])
-
-        def fetch(self, message_set, message_parts):
-            return (
-                'OK', [
-                    (
-                        '1 (RFC822 {4800}',
-                        TEST_EMAIL_BASE64_FILENAME
-                    ), ' FLAGS (\\Seen))'
-                ]
-            )
-
-        def store(self, message_set, command, flags):
-            return ('OK', ['1 (FLAGS (\\Seen \\Deleted))'])
-
-        def expunge(self):
-            return ('OK', ['1'])
-
-        def close(self):
-            return ('OK', ['Returned to authenticated state. (Success)'])
-
-        def logout(self):
-            return ('BYE', ['LOGOUT Requested'])
-
    @mock.patch('imaplib.IMAP4_SSL', autospec=True)
    def test_download_document(self, mock_imaplib):
-        mock_imaplib.return_value = IMAPSourceTestCase.MockIMAPServer()
+        mock_imaplib.return_value = MockIMAPServer()
        self.source = IMAPEmail.objects.create(
            document_type=self.test_document_type, label='', host='',
            password='', username=''