From 4facc93670d138f0c4cb0124a9b4ca8fe32119b4 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@protonmail.com>
Date: Thu, 30 May 2019 15:24:18 +0200
Subject: [PATCH] retention: Add archiving of SubMessages.

---
 zerver/lib/retention.py        |  9 ++++++-
 zerver/tests/test_retention.py | 43 +++++++++++++++++++++++++++++++---
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/zerver/lib/retention.py b/zerver/lib/retention.py
index 1cfe30a154..e8334315fc 100644
--- a/zerver/lib/retention.py
+++ b/zerver/lib/retention.py
@@ -4,7 +4,8 @@ from datetime import timedelta
 from django.db import connection, transaction
 from django.utils.timezone import now as timezone_now
 from zerver.models import (Message, UserMessage, ArchivedMessage, ArchivedUserMessage, Realm,
-                           Attachment, ArchivedAttachment, Reaction, ArchivedReaction)
+                           Attachment, ArchivedAttachment, Reaction, ArchivedReaction,
+                           SubMessage, ArchivedSubMessage)
 
 from typing import Any, Dict, List
 
@@ -15,6 +16,12 @@ models_with_message_key = [
         'table_name': 'zerver_reaction',
         'archive_table_name': 'zerver_archivedreaction'
     },
+    {
+        'class': SubMessage,
+        'archive_class': ArchivedSubMessage,
+        'table_name': 'zerver_submessage',
+        'archive_table_name': 'zerver_archivedsubmessage'
+    },
 ]  # type: List[Dict[str, Any]]
 
 @transaction.atomic
diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py
index cb8bf1e61b..b27dcbe4bc 100644
--- a/zerver/tests/test_retention.py
+++ b/zerver/tests/test_retention.py
@@ -9,7 +9,7 @@ from zerver.lib.test_classes import ZulipTestCase
 from zerver.lib.upload import create_attachment
 from zerver.models import (Message, Realm, UserProfile, ArchivedUserMessage, SubMessage,
                            ArchivedMessage, Attachment, ArchivedAttachment, UserMessage,
-                           Reaction, ArchivedReaction,
+                           Reaction, ArchivedReaction, ArchivedSubMessage,
                            get_realm, get_user_profile_by_email, get_system_bot)
 from zerver.lib.retention import (
     archive_messages,
@@ -307,8 +307,6 @@ class TestArchivingGeneral(RetentionTestingBase):
 
 class TestArchivingSubMessages(RetentionTestingBase):
     def test_archiving_submessages(self) -> None:
-        # TODO: Expand this accordingly, when archiving submessages is actually implemented.
-        # For now, we just test if submessages of an archived message get correctly deleted.
         expired_msg_ids = self._make_expired_zulip_messages(2)
         cordelia = self.example_user('cordelia')
         hamlet = self.example_user('hamlet')
@@ -345,6 +343,11 @@ class TestArchivingSubMessages(RetentionTestingBase):
         archive_messages()
         self.assertEqual(SubMessage.objects.filter(id__in=submessage_ids).count(), 0)
 
+        self.assertEqual(
+            set(ArchivedSubMessage.objects.filter(id__in=submessage_ids).values_list('id', flat=True)),
+            set(submessage_ids)
+        )
+
 class TestArchivingReactions(RetentionTestingBase, EmojiReactionBase):
     def test_archiving_reactions(self) -> None:
         expired_msg_ids = self._make_expired_zulip_messages(2)
@@ -525,6 +528,40 @@ class MoveMessageToArchiveGeneral(MoveMessageToArchiveBase):
         # Now the attachment should have been deleted:
         self.assertEqual(Attachment.objects.count(), 0)
 
+class MoveMessageToArchiveWithSubMessages(MoveMessageToArchiveBase):
+    def test_archiving_message_with_submessages(self) -> None:
+        msg_id = self.send_stream_message(self.sender, "Verona")
+        cordelia = self.example_user('cordelia')
+        hamlet = self.example_user('hamlet')
+
+        do_add_submessage(
+            realm=get_realm('zulip'),
+            sender_id=cordelia.id,
+            message_id=msg_id,
+            msg_type='whatever',
+            content='{"name": "alice", "salary": 20}'
+        )
+        do_add_submessage(
+            realm=get_realm('zulip'),
+            sender_id=hamlet.id,
+            message_id=msg_id,
+            msg_type='whatever',
+            content='{"name": "john", "salary": 30}'
+        )
+
+        submessage_ids = list(
+            SubMessage.objects.filter(message_id=msg_id).values_list('id', flat=True)
+        )
+
+        self.assertEqual(SubMessage.objects.filter(id__in=submessage_ids).count(), 2)
+        move_messages_to_archive(message_ids=[msg_id])
+
+        self.assertEqual(
+            set(ArchivedSubMessage.objects.filter(message_id=msg_id).values_list("id", flat=True)),
+            set(submessage_ids)
+        )
+        self.assertEqual(SubMessage.objects.filter(id__in=submessage_ids).count(), 0)
+
 class MoveMessageToArchiveWithReactions(MoveMessageToArchiveBase, EmojiReactionBase):
     def test_archiving_message_with_reactions(self) -> None:
         msg_id = self.send_stream_message(self.sender, "Verona")