export: Add support for exporting realm with member consent.

This lets us handle directly in our tooling the user experience that
we document for exporting a realm with member consent (before, it
required unpleasant manual work).

templates/zerver/help/export-your-organization.md

@@ -51,7 +51,12 @@ Such an export will include all the messages received by any user in
 the organization that consented to the data export.  In particular, it
 will include all public stream content and any private stream or
 private message content where at least one of the participants gave
-consent.
+consent.  Users who do not consent to export their private data will
+still be able to see public message history sent before the export,
+but will not have access to any private message history, including the
+history of which messages they personally received (which will result
+in their experience for All Messages and our full-text search being
+like that of a new user who joined after the data export).
 
 For **full export without member consent**, we will additionally need
 evidence that you have authority to read members' private

tools/linter_lib/custom_check.py

@@ -559,7 +559,9 @@ def build_custom_checkers(by_lang):
          'exclude': set(["zerver/tests",
                          "zerver/lib/onboarding.py",
                          "zilencer/management/commands/add_mock_conversation.py",
-                         "zerver/worker/queue_processors.py"]),
+                         "zerver/worker/queue_processors.py",
+                         "zerver/management/commands/export.py",
+                         "zerver/lib/export.py"]),
          'description': 'Please use access_message() to fetch Message objects',
          },
         {'pattern': 'Stream.objects.get',

zerver/lib/export.py

@@ -15,6 +15,7 @@ from django.forms.models import model_to_dict
 from django.utils.timezone import make_aware as timezone_make_aware
 from django.utils.timezone import is_naive as timezone_is_naive
 from django.core.management.base import CommandError
+from django.db.models import Q
 import glob
 import logging
 import os
@@ -907,11 +908,15 @@ def fetch_huddle_objects(response: TableData, config: Config, context: Context)
 def fetch_usermessages(realm: Realm,
                        message_ids: Set[int],
                        user_profile_ids: Set[int],
-                       message_filename: Path) -> List[Record]:
+                       message_filename: Path,
+                       consent_message_id: Optional[int]=None) -> List[Record]:
     # UserMessage export security rule: You can export UserMessages
     # for the messages you exported for the users in your realm.
     user_message_query = UserMessage.objects.filter(user_profile__realm=realm,
                                                     message_id__in=message_ids)
+    if consent_message_id is not None:
+        consented_user_ids = get_consented_user_ids(consent_message_id)
+        user_profile_ids = user_profile_ids & consented_user_ids
     user_message_chunk = []
     for user_message in user_message_query:
         if user_message.user_profile_id not in user_profile_ids:
@@ -923,7 +928,8 @@ def fetch_usermessages(realm: Realm,
     logging.info("Fetched UserMessages for %s" % (message_filename,))
     return user_message_chunk
 
-def export_usermessages_batch(input_path: Path, output_path: Path) -> None:
+def export_usermessages_batch(input_path: Path, output_path: Path,
+                              consent_message_id: Optional[int]=None) -> None:
     """As part of the system for doing parallel exports, this runs on one
     batch of Message objects and adds the corresponding UserMessage
     objects. (This is called by the export_usermessage_batch
@@ -935,7 +941,8 @@ def export_usermessages_batch(input_path: Path, output_path: Path) -> None:
     del output['zerver_userprofile_ids']
     realm = Realm.objects.get(id=output['realm_id'])
     del output['realm_id']
-    output['zerver_usermessage'] = fetch_usermessages(realm, set(message_ids), user_profile_ids, output_path)
+    output['zerver_usermessage'] = fetch_usermessages(realm, set(message_ids), user_profile_ids,
+                                                      output_path, consent_message_id)
     write_message_export(output_path, output)
     os.unlink(input_path)
 
@@ -947,7 +954,8 @@ def export_partial_message_files(realm: Realm,
                                  response: TableData,
                                  chunk_size: int=MESSAGE_BATCH_CHUNK_SIZE,
                                  output_dir: Optional[Path]=None,
-                                 public_only: bool=False) -> Set[int]:
+                                 public_only: bool=False,
+                                 consent_message_id: Optional[int]=None) -> Set[int]:
     if output_dir is None:
         output_dir = tempfile.mkdtemp(prefix="zulip-export")
 
@@ -976,32 +984,57 @@ def export_partial_message_files(realm: Realm,
         response['zerver_userprofile_mirrordummy'] +
         response['zerver_userprofile_crossrealm'])
 
+    consented_user_ids = set()  # type: Set[int]
+    if consent_message_id is not None:
+        consented_user_ids = get_consented_user_ids(consent_message_id)
+
     if public_only:
         recipient_streams = Stream.objects.filter(realm=realm, invite_only=False)
         recipient_ids = Recipient.objects.filter(
             type=Recipient.STREAM, type_id__in=recipient_streams).values_list("id", flat=True)
         recipient_ids_for_us = get_ids(response['zerver_recipient']) & set(recipient_ids)
+    elif consent_message_id is not None:
+        public_streams = Stream.objects.filter(realm=realm, invite_only=False)
+        public_stream_recipient_ids = Recipient.objects.filter(
+            type=Recipient.STREAM, type_id__in=public_streams).values_list("id", flat=True)
+
+        consented_recipient_ids = Subscription.objects.filter(user_profile__id__in=consented_user_ids). \
+            values_list("recipient_id", flat=True)
+
+        recipient_ids = set(public_stream_recipient_ids) | set(consented_recipient_ids)
+        recipient_ids_for_us = get_ids(response['zerver_recipient']) & recipient_ids
     else:
         recipient_ids_for_us = get_ids(response['zerver_recipient'])
-
-    # We capture most messages here, since the
-    # recipients we subscribe to are also the
-    # recipients of most messages we send.
-    messages_we_received = Message.objects.filter(
-        sender__in=ids_of_our_possible_senders,
-        recipient__in=recipient_ids_for_us,
-    ).order_by('id')
+        # For a full export, we have implicit consent for all users in the export.
+        consented_user_ids = user_ids_for_us
 
     if public_only:
+        messages_we_received = Message.objects.filter(
+            sender__in=ids_of_our_possible_senders,
+            recipient__in=recipient_ids_for_us,
+        ).order_by('id')
+
         # For the public stream export, we only need the messages those streams received.
         message_queries = [
             messages_we_received,
         ]
     else:
-        # This should pick up stragglers; messages we sent
-        # where we the recipient wasn't subscribed to by any of
-        # us (such as PMs to "them").
-        ids_of_non_exported_possible_recipients = ids_of_our_possible_senders - user_ids_for_us
+        # We capture most messages here: Messages that were sent by
+        # anyone in the export and received by any of the users who we
+        # have consent to export.
+        messages_we_received = Message.objects.filter(
+            sender__in=ids_of_our_possible_senders,
+            recipient__in=recipient_ids_for_us,
+        ).order_by('id')
+
+        # The above query is missing some messages that consenting
+        # users have access to, namely, PMs sent by one of the users
+        # in our export to another user (since the only subscriber to
+        # a Recipient object for Recipient.PERSONAL is the recipient,
+        # not the sender).  The `consented_user_ids` list has
+        # precisely those users whose Recipient.PERSONAL recipient ID
+        # was already present in recipient_ids_for_us above.
+        ids_of_non_exported_possible_recipients = ids_of_our_possible_senders - consented_user_ids
 
         recipients_for_them = Recipient.objects.filter(
             type=Recipient.PERSONAL,
@@ -1009,10 +1042,14 @@ def export_partial_message_files(realm: Realm,
         recipient_ids_for_them = get_ids(recipients_for_them)
 
         messages_we_sent_to_them = Message.objects.filter(
-            sender__in=user_ids_for_us,
+            sender__in=consented_user_ids,
             recipient__in=recipient_ids_for_them,
         ).order_by('id')
 
+        messages_we_received = Message.objects.filter(
+            Q(sender__in=consented_user_ids) | Q(recipient__in=recipient_ids_for_us),
+        )
+
         message_queries = [
             messages_we_received,
             messages_we_sent_to_them,
@@ -1357,7 +1394,8 @@ def do_write_stats_file_for_realm_export(output_dir: Path) -> None:
 
 def do_export_realm(realm: Realm, output_dir: Path, threads: int,
                     exportable_user_ids: Optional[Set[int]]=None,
-                    public_only: bool=False) -> None:
+                    public_only: bool=False,
+                    consent_message_id: Optional[int]=None) -> None:
     response = {}  # type: TableData
 
     # We need at least one thread running to export
@@ -1391,7 +1429,8 @@ def do_export_realm(realm: Realm, output_dir: Path, threads: int,
     # have millions of messages.
     logging.info("Exporting .partial files messages")
     message_ids = export_partial_message_files(realm, response, output_dir=output_dir,
-                                               public_only=public_only)
+                                               public_only=public_only,
+                                               consent_message_id=consent_message_id)
     logging.info('%d messages were exported' % (len(message_ids),))
 
     # zerver_reaction
@@ -1411,7 +1450,8 @@ def do_export_realm(realm: Realm, output_dir: Path, threads: int,
     export_attachment_table(realm=realm, output_dir=output_dir, message_ids=message_ids)
 
     # Start parallel jobs to export the UserMessage objects.
-    launch_user_message_subprocesses(threads=threads, output_dir=output_dir)
+    launch_user_message_subprocesses(threads=threads, output_dir=output_dir,
+                                     consent_message_id=consent_message_id)
 
     logging.info("Finished exporting %s" % (realm.string_id,))
     create_soft_link(source=output_dir, in_progress=False)
@@ -1445,14 +1485,21 @@ def create_soft_link(source: Path, in_progress: bool=True) -> None:
     if is_done:
         logging.info('See %s for output files' % (new_target,))
 
-
-def launch_user_message_subprocesses(threads: int, output_dir: Path) -> None:
+def launch_user_message_subprocesses(threads: int, output_dir: Path,
+                                     consent_message_id: Optional[int]=None) -> None:
     logging.info('Launching %d PARALLEL subprocesses to export UserMessage rows' % (threads,))
 
     def run_job(shard: str) -> int:
-        subprocess.call([os.path.join(settings.DEPLOY_ROOT, "manage.py"),
-                         'export_usermessage_batch', '--path',
-                         str(output_dir), '--thread', shard])
+        arguments = [
+            os.path.join(settings.DEPLOY_ROOT, "manage.py"),
+            'export_usermessage_batch',
+            '--path', str(output_dir),
+            '--thread', shard
+        ]
+        if consent_message_id is not None:
+            arguments.extend(['--consent-message-id', str(consent_message_id)])
+
+        subprocess.call(arguments)
         return 0
 
     for (status, job) in run_parallel(run_job,
@@ -1507,6 +1554,16 @@ def get_single_user_config() -> Config:
     )
 
     # zerver_stream
+    #
+    # TODO: We currently export the existence of private streams, but
+    # not their message history, in the "export with partial member
+    # consent" code path.  This consistent with our documented policy,
+    # since that data is available to the organization administrator
+    # who initiated the export, but unnecessary and potentially
+    # confusing; it'd be better to just skip those streams from the
+    # export (which would require more complex export logic for the
+    # subscription/recipient/stream tables to exclude private streams
+    # with no consenting subscribers).
     Config(
         table='zerver_stream',
         model=Stream,
@@ -1598,11 +1655,20 @@ def get_analytics_config() -> Config:
 
     return analytics_config
 
+def get_consented_user_ids(consent_message_id: int) -> Set[int]:
+    return set(Reaction.objects.filter(message__id=consent_message_id,
+                                       reaction_type="unicode_emoji",
+                                       # thumbsup = 1f44d
+                                       emoji_code="1f44d").
+               values_list("user_profile", flat=True))
 
 def export_realm_wrapper(realm: Realm, output_dir: str,
                          threads: int, upload_to_s3: bool,
-                         public_only: bool, delete_after_upload: bool) -> Optional[str]:
-    do_export_realm(realm=realm, output_dir=output_dir, threads=threads, public_only=public_only)
+                         public_only: bool,
+                         delete_after_upload: bool,
+                         consent_message_id: Optional[int]=None) -> Optional[str]:
+    do_export_realm(realm=realm, output_dir=output_dir, threads=threads,
+                    public_only=public_only, consent_message_id=consent_message_id)
     print("Finished exporting to %s; tarring" % (output_dir,))
 
     do_write_stats_file_for_realm_export(output_dir)

zerver/management/commands/export.py

@@ -9,6 +9,7 @@ from django.core.management.base import CommandError
 
 from zerver.lib.management import ZulipBaseCommand
 from zerver.lib.export import export_realm_wrapper
+from zerver.models import Message, Reaction
 
 class Command(ZulipBaseCommand):
     help = """Exports all data from a Zulip realm
@@ -92,6 +93,12 @@ class Command(ZulipBaseCommand):
         parser.add_argument('--public-only',
                             action="store_true",
                             help='Export only public stream messages and associated attachments')
+        parser.add_argument('--consent-message-id',
+                            dest="consent_message_id",
+                            action="store",
+                            default=None,
+                            type=int,
+                            help='ID of the message advertising users to react with thumbs up')
         parser.add_argument('--upload-to-s3',
                             action="store_true",
                             help="Whether to upload resulting tarball to s3")
@@ -105,6 +112,9 @@ class Command(ZulipBaseCommand):
         assert realm is not None  # Should be ensured by parser
 
         output_dir = options["output_dir"]
+        public_only = options["public_only"]
+        consent_message_id = options["consent_message_id"]
+
         if output_dir is None:
             output_dir = tempfile.mkdtemp(prefix="zulip-export-")
         else:
@@ -112,13 +122,42 @@ class Command(ZulipBaseCommand):
         if os.path.exists(output_dir):
             shutil.rmtree(output_dir)
         os.makedirs(output_dir)
-        print("Exporting realm %s" % (realm.string_id,))
+        print("\033[94mExporting realm\033[0m: %s" % (realm.string_id,))
+
         num_threads = int(options['threads'])
         if num_threads < 1:
             raise CommandError('You must have at least one thread.')
 
+        if public_only and consent_message_id is not None:
+            raise CommandError('Please pass either --public-only or --consent-message-id')
+
+        if consent_message_id is not None:
+            try:
+                message = Message.objects.get(id=consent_message_id)
+            except Message.DoesNotExist:
+                raise CommandError("Message with given ID does not exist. Aborting...")
+
+            if message.last_edit_time is not None:
+                raise CommandError("Message was edited. Aborting...")
+
+            # Since the message might have been sent by
+            # Notification Bot, we can't trivially check the realm of
+            # the message through message.sender.realm.  So instead we
+            # check the realm of the people who reacted to the message
+            # (who must all be in the message's realm).
+            reactions = Reaction.objects.filter(message=message, emoji_code="1f44d",
+                                                reaction_type="unicode_emoji")
+            for reaction in reactions:
+                if reaction.user_profile.realm != realm:
+                    raise CommandError("Users from a different realm reacted to message. Aborting...")
+
+            print("\n\033[94mMessage content:\033[0m\n{}\n".format(message.content))
+
+            print("\033[94mNumber of users that reacted +1:\033[0m {}\n".format(len(reactions)))
+
         # Allows us to trigger exports separately from command line argument parsing
         export_realm_wrapper(realm=realm, output_dir=output_dir,
                              threads=num_threads, upload_to_s3=options['upload_to_s3'],
-                             public_only=options["public_only"],
-                             delete_after_upload=options["delete_after_upload"])
+                             public_only=public_only,
+                             delete_after_upload=options["delete_after_upload"],
+                             consent_message_id=consent_message_id)

zerver/management/commands/export_usermessage_batch.py

@@ -24,6 +24,12 @@ class Command(BaseCommand):
                             action="store",
                             default=None,
                             help='Thread ID')
+        parser.add_argument('--consent-message-id',
+                            dest="consent_message_id",
+                            action="store",
+                            default=None,
+                            type=int,
+                            help='ID of the message advertising users to react with thumbs up')
 
     def handle(self, *args: Any, **options: Any) -> None:
         logging.info("Starting UserMessage batch thread %s" % (options['thread'],))
@@ -38,7 +44,7 @@ class Command(BaseCommand):
                 continue
             logging.info("Thread %s processing %s" % (options['thread'], output_path))
             try:
-                export_usermessages_batch(locked_path, output_path)
+                export_usermessages_batch(locked_path, output_path, options["consent_message_id"])
             except Exception:
                 # Put the item back in the free pool when we fail
                 shutil.move(locked_path, partial_path)

zerver/tests/test_import_export.py

@@ -8,6 +8,7 @@ import ujson
 from mock import patch
 from typing import Any, Dict, List, Set, Optional, Tuple, Callable, \
     FrozenSet
+from django.db.models import Q
 
 from zerver.lib.export import (
     do_export_realm,
@@ -49,6 +50,8 @@ from zerver.lib.bot_config import (
 )
 from zerver.lib.actions import (
     do_create_user,
+    do_add_reaction,
+    create_stream_if_needed
 )
 
 from zerver.lib.test_runner import slow
@@ -61,6 +64,7 @@ from zerver.models import (
     Subscription,
     Attachment,
     RealmEmoji,
+    Reaction,
     Recipient,
     UserMessage,
     CustomProfileField,
@@ -224,7 +228,8 @@ class ImportExportTest(ZulipTestCase):
         os.makedirs(output_dir, exist_ok=True)
         return output_dir
 
-    def _export_realm(self, realm: Realm, exportable_user_ids: Optional[Set[int]]=None) -> Dict[str, Any]:
+    def _export_realm(self, realm: Realm, exportable_user_ids: Optional[Set[int]]=None,
+                      consent_message_id: Optional[int]=None) -> Dict[str, Any]:
         output_dir = self._make_output_dir()
         with patch('logging.info'), patch('zerver.lib.export.create_soft_link'):
             do_export_realm(
@@ -232,12 +237,14 @@ class ImportExportTest(ZulipTestCase):
                 output_dir=output_dir,
                 threads=0,
                 exportable_user_ids=exportable_user_ids,
+                consent_message_id=consent_message_id,
             )
             # 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, 'messages-000001.json')
+                output_path=os.path.join(output_dir, 'messages-000001.json'),
+                consent_message_id=consent_message_id,
             )
 
         def read_file(fn: str) -> Any:
@@ -442,6 +449,120 @@ class ImportExportTest(ZulipTestCase):
         self.assertIn(self.example_email('iago'), dummy_user_emails)
         self.assertNotIn(self.example_email('cordelia'), dummy_user_emails)
 
+    def test_export_realm_with_member_consent(self) -> None:
+        realm = Realm.objects.get(string_id='zulip')
+
+        # Create private streams and subscribe users for testing export
+        create_stream_if_needed(realm, "Private A", invite_only=True)
+        self.subscribe(self.example_user("iago"), "Private A")
+        self.subscribe(self.example_user("othello"), "Private A")
+        self.send_stream_message(self.example_email("iago"), "Private A", "Hello Stream A")
+
+        create_stream_if_needed(realm, "Private B", invite_only=True)
+        self.subscribe(self.example_user("hamlet"), "Private B")
+        self.subscribe(self.example_user("prospero"), "Private B")
+        self.send_stream_message(self.example_email("prospero"), "Private B", "Hello Stream B")
+
+        create_stream_if_needed(realm, "Private C", invite_only=True)
+        self.subscribe(self.example_user("othello"), "Private C")
+        self.subscribe(self.example_user("prospero"), "Private C")
+        stream_c_message_id = self.send_stream_message(self.example_email("othello"),
+                                                       "Private C", "Hello Stream C")
+
+        # Create huddles
+        self.send_huddle_message(self.example_email("iago"), [self.example_email("cordelia"),
+                                                              self.example_email("AARON")])
+        huddle_a = Huddle.objects.last()
+        self.send_huddle_message(self.example_email("ZOE"), [self.example_email("hamlet"),
+                                                             self.example_email("AARON"),
+                                                             self.example_email("othello")])
+        huddle_b = Huddle.objects.last()
+
+        huddle_b_message_id = self.send_huddle_message(
+            self.example_email("AARON"), [self.example_email("cordelia"),
+                                          self.example_email("ZOE"),
+                                          self.example_email("othello")])
+
+        # Send message advertising export and make users react
+        self.send_stream_message(self.example_email("othello"), "Verona",
+                                 topic_name="Export",
+                                 content="Thumbs up for export")
+        message = Message.objects.last()
+        consented_user_ids = [self.example_user(user).id for user in ["iago", "hamlet"]]
+        do_add_reaction(self.example_user("iago"), message, "+1", "1f44d",  Reaction.UNICODE_EMOJI)
+        do_add_reaction(self.example_user("hamlet"), message, "+1", "1f44d",  Reaction.UNICODE_EMOJI)
+
+        realm_emoji = RealmEmoji.objects.get(realm=realm)
+        realm_emoji.delete()
+        full_data = self._export_realm(realm, consent_message_id=message.id)
+        realm_emoji.save()
+
+        data = full_data['realm']
+
+        self.assertEqual(len(data['zerver_userprofile_crossrealm']), 0)
+        self.assertEqual(len(data['zerver_userprofile_mirrordummy']), 0)
+
+        def get_set(table: str, field: str) -> Set[str]:
+            values = set(r[field] for r in data[table])
+            return values
+
+        def find_by_id(table: str, db_id: int) -> 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(self.example_email('cordelia'), exported_user_emails)
+        self.assertIn(self.example_email('hamlet'), exported_user_emails)
+        self.assertIn(self.example_email('iago'), exported_user_emails)
+        self.assertIn(self.example_email('othello'), 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',
+                 u'Private A', u'Private B', u'Private C'])
+        )
+
+        data = full_data['message']
+        exported_usermessages = UserMessage.objects.filter(user_profile__in=[self.example_user("iago"),
+                                                                             self.example_user("hamlet")])
+        um = exported_usermessages[0]
+        self.assertEqual(len(data["zerver_usermessage"]), len(exported_usermessages))
+        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)
+
+        public_stream_names = ['Denmark', 'Rome', 'Scotland', 'Venice', 'Verona']
+        public_stream_ids = Stream.objects.filter(name__in=public_stream_names).values_list("id", flat=True)
+        public_stream_recipients = Recipient.objects.filter(type_id__in=public_stream_ids, type=Recipient.STREAM)
+        public_stream_message_ids = Message.objects.filter(recipient__in=public_stream_recipients).values_list("id", flat=True)
+
+        # Messages from Private Stream C are not exported since no member gave consent
+        private_stream_ids = Stream.objects.filter(name__in=["Private A", "Private B"]).values_list("id", flat=True)
+        private_stream_recipients = Recipient.objects.filter(type_id__in=private_stream_ids, type=Recipient.STREAM)
+        private_stream_message_ids = Message.objects.filter(recipient__in=private_stream_recipients).values_list("id", flat=True)
+
+        pm_recipients = Recipient.objects.filter(type_id__in=consented_user_ids, type=Recipient.PERSONAL)
+        pm_query = Q(recipient__in=pm_recipients) | Q(sender__in=consented_user_ids)
+        exported_pm_ids = Message.objects.filter(pm_query).values_list("id", flat=True).values_list("id", flat=True)
+
+        # Third huddle is not exported since none of the members gave consent
+        huddle_recipients = Recipient.objects.filter(type_id__in=[huddle_a.id, huddle_b.id], type=Recipient.HUDDLE)
+        pm_query = Q(recipient__in=huddle_recipients) | Q(sender__in=consented_user_ids)
+        exported_huddle_ids = Message.objects.filter(pm_query).values_list("id", flat=True).values_list("id", flat=True)
+
+        exported_msg_ids = set(public_stream_message_ids) | set(private_stream_message_ids) \
+            | set(exported_pm_ids) | set(exported_huddle_ids)
+        self.assertEqual(get_set("zerver_message", "id"), exported_msg_ids)
+        self.assertNotIn(stream_c_message_id, exported_msg_ids)
+        self.assertNotIn(huddle_b_message_id, exported_msg_ids)
+
     def test_export_single_user(self) -> None:
         output_dir = self._make_output_dir()
         cordelia = self.example_user('cordelia')

zerver/tests/test_management_commands.py

@@ -5,22 +5,24 @@ import os
 import re
 from datetime import timedelta
 from email.utils import parseaddr
+import mock
 from mock import MagicMock, patch, call
 from typing import List, Dict, Any, Optional
 
 from django.conf import settings
 from django.core.management import call_command
 from django.test import TestCase, override_settings
-from zerver.lib.actions import do_create_user
+from zerver.lib.actions import do_create_user, do_add_reaction
 from zerver.lib.management import ZulipBaseCommand, CommandError, check_config
 from zerver.lib.test_classes import ZulipTestCase
 from zerver.lib.test_helpers import stdout_suppressed
 from zerver.lib.test_runner import slow
 from zerver.models import Recipient, get_user_profile_by_email, get_stream
+from django.utils.timezone import now as timezone_now
 
 from zerver.management.commands.send_webhook_fixture_message import Command
 from zerver.lib.test_helpers import most_recent_message
-from zerver.models import get_realm, UserProfile, Realm
+from zerver.models import get_realm, UserProfile, Realm, Reaction, Message
 from confirmation.models import RealmCreationKey, generate_realm_creation_url
 
 class TestCheckConfig(ZulipTestCase):
@@ -400,3 +402,35 @@ class TestInvoicePlans(ZulipTestCase):
             call_command(self.COMMAND_NAME)
 
         m.assert_called_once()
+
+class TestExport(ZulipTestCase):
+    COMMAND_NAME = 'export'
+
+    def test_command_with_consented_message_id(self) -> None:
+        realm = get_realm("zulip")
+        self.send_stream_message(self.example_email("othello"), "Verona",
+                                 topic_name="Export",
+                                 content="Thumbs up for export")
+        message = Message.objects.last()
+        do_add_reaction(self.example_user("iago"), message, "+1", "1f44d",  Reaction.UNICODE_EMOJI)
+        do_add_reaction(self.example_user("hamlet"), message, "+1", "1f44d",  Reaction.UNICODE_EMOJI)
+
+        with patch("zerver.management.commands.export.export_realm_wrapper") as m:
+            call_command(self.COMMAND_NAME, "-r=zulip", "--consent-message-id={}".format(message.id))
+            m.assert_called_once_with(realm=realm, public_only=False, consent_message_id=message.id,
+                                      delete_after_upload=False, threads=mock.ANY, output_dir=mock.ANY,
+                                      upload_to_s3=False)
+
+        with self.assertRaisesRegex(CommandError, "Message with given ID does not"):
+            call_command(self.COMMAND_NAME, "-r=zulip", "--consent-message-id=123456")
+
+        message.last_edit_time = timezone_now()
+        message.save()
+        with self.assertRaisesRegex(CommandError, "Message was edited. Aborting..."):
+            call_command(self.COMMAND_NAME, "-r=zulip", "--consent-message-id={}".format(message.id))
+
+        message.last_edit_time = None
+        message.save()
+        do_add_reaction(self.mit_user("sipbtest"), message, "+1", "1f44d",  Reaction.UNICODE_EMOJI)
+        with self.assertRaisesRegex(CommandError, "Users from a different realm reacted to message. Aborting..."):
+            call_command(self.COMMAND_NAME, "-r=zulip", "--consent-message-id={}".format(message.id))