export: Add basic export tests. (fixes #1584)

zerver/lib/export.py

@@ -1006,6 +1006,7 @@ def do_export_realm(realm, output_dir, threads, exportable_user_ids=None):
     # We need at least one thread running to export
     # UserMessage rows.  The management command should
     # enforce this for us.
+    if not settings.TEST_SUITE:
         assert threads >= 1
 
     assert os.path.exists("./manage.py")

zerver/tests/test_export.py

@@ -0,0 +1,172 @@
+# -*- coding: utf-8 -*-
+from __future__ import absolute_import
+from __future__ import print_function
+
+from django.conf import settings
+from django.test import TestCase
+
+import os
+import shutil
+import ujson
+
+from mock import patch, MagicMock
+from typing import Any
+
+from zerver.lib.actions import (
+    do_claim_attachments,
+)
+
+from zerver.lib.export import (
+    do_export_realm,
+    export_usermessages_batch,
+)
+from zerver.lib.upload import (
+    claim_attachment,
+    upload_message_image,
+)
+from zerver.lib.utils import mkdir_p
+from zerver.models import (
+    get_user_profile_by_email,
+    Message,
+    Realm,
+    UserMessage,
+)
+
+def rm_tree(path):
+    # type: (str) -> None
+    if os.path.exists(path):
+        shutil.rmtree(path)
+
+class ExportTest(TestCase):
+
+    def setUp(self):
+        # type: () -> None
+        rm_tree(settings.LOCAL_UPLOADS_DIR)
+
+    def _make_output_dir(self):
+        # type: () -> str
+        output_dir = 'var/test-export'
+        rm_tree(output_dir)
+        mkdir_p(output_dir)
+        return output_dir
+
+    def _export_realm(self, domain, exportable_user_ids=None):
+        # type: (str, Set[int]) -> Dict[str, Any]
+        output_dir = self._make_output_dir()
+        realm = Realm.objects.get(domain=domain)
+        with patch('logging.info'), patch('zerver.lib.export.create_soft_link'):
+            do_export_realm(
+                realm=realm,
+                output_dir=output_dir,
+                threads=0,
+                exportable_user_ids=exportable_user_ids,
+            )
+            # TODO: Process the second partial file, which can be created
+            #       for certain edge cases.
+            export_usermessages_batch(
+                input_path=os.path.join(output_dir, 'messages-000001.json.partial'),
+                output_path=os.path.join(output_dir, 'message.json')
+            )
+
+        def read_file(fn):
+            full_fn = os.path.join(output_dir, fn)
+            with open(full_fn) as f:
+                return ujson.load(f)
+
+        result = {}
+        result['realm'] = read_file('realm.json')
+        result['attachment'] = read_file('attachment.json')
+        result['message'] = read_file('message.json')
+        result['uploads_dir'] = os.path.join(output_dir, 'uploads')
+        return result
+
+    def test_attachment(self):
+        # type: () -> None
+        message = Message.objects.all()[0]
+        user_profile = message.sender
+        url = upload_message_image(u'dummy.txt', u'text/plain', b'zulip!', user_profile)
+        path_id = url.replace('/user_uploads/', '')
+        claim_attachment(
+            user_profile=user_profile,
+            path_id=path_id,
+            message=message,
+            is_message_realm_public=True
+        )
+
+        domain = 'zulip.com'
+        full_data = self._export_realm(domain=domain)
+
+        data = full_data['attachment']
+        self.assertEqual(len(data['zerver_attachment']), 1)
+        record = data['zerver_attachment'][0]
+        self.assertEqual(record['path_id'], path_id)
+
+        fn = os.path.join(full_data['uploads_dir'], path_id)
+        with open(fn) as f:
+            self.assertEqual(f.read(), 'zulip!')
+
+    def test_zulip_realm(self):
+        # type: () -> None
+        domain = 'zulip.com'
+        full_data = self._export_realm(domain=domain)
+
+        data = full_data['realm']
+        self.assertEqual(len(data['zerver_userprofile_crossrealm']), 0)
+        self.assertEqual(len(data['zerver_userprofile_mirrordummy']), 0)
+
+        def get_set(table, field):
+            # type: (str, str) -> Set[str]
+            values = set(r[field] for r in data[table])
+            # print('set(%s)' % sorted(values))
+            return values
+
+        def find_by_id(table, db_id):
+            # type: (str) -> Dict[str, Any]
+            return [
+                r for r in data[table]
+                if r['id'] == db_id][0]
+
+
+        exported_user_emails = get_set('zerver_userprofile', 'email')
+        self.assertIn('cordelia@zulip.com', exported_user_emails)
+        self.assertIn('default-bot@zulip.com', exported_user_emails)
+        self.assertIn('emailgateway@zulip.com', exported_user_emails)
+
+        exported_streams = get_set('zerver_stream', 'name')
+        self.assertEqual(
+            exported_streams,
+            set([u'Denmark', u'Rome', u'Scotland', u'Venice', u'Verona'])
+        )
+
+        data = full_data['message']
+        um = UserMessage.objects.all()[0]
+        exported_um = find_by_id('zerver_usermessage', um.id)
+        self.assertEqual(exported_um['message'], um.message_id)
+        self.assertEqual(exported_um['user_profile'], um.user_profile_id)
+
+        exported_message = find_by_id('zerver_message', um.message_id)
+        self.assertEqual(exported_message['content'], um.message.content)
+
+
+        # TODO, extract get_set/find_by_id, so we can split this test up
+
+        # Now, restrict users
+        cordelia = get_user_profile_by_email('cordelia@zulip.com')
+        hamlet = get_user_profile_by_email('hamlet@zulip.com')
+        user_ids = set([cordelia.id, hamlet.id])
+
+        full_data = self._export_realm(
+            domain=domain,
+            exportable_user_ids=user_ids
+        )
+        data = full_data['realm']
+        exported_user_emails = get_set('zerver_userprofile', 'email')
+        self.assertIn('cordelia@zulip.com', exported_user_emails)
+        self.assertIn('hamlet@zulip.com', exported_user_emails)
+        self.assertNotIn('default-bot@zulip.com', exported_user_emails)
+        self.assertNotIn('iago@zulip.com', exported_user_emails)
+
+        dummy_user_emails = get_set('zerver_userprofile_mirrordummy', 'email')
+        self.assertIn('iago@zulip.com', dummy_user_emails)
+        self.assertNotIn('cordelia@zulip.com', dummy_user_emails)
+