From 29729b77488d2804154b044d177278d6442d40d6 Mon Sep 17 00:00:00 2001
From: Mateusz Mandera <mateusz.mandera@protonmail.com>
Date: Thu, 30 May 2019 12:49:20 +0200
Subject: [PATCH] test_retention: Check that SubMessages get correctly deleted.

---
 zerver/tests/test_retention.py | 44 ++++++++++++++++++++++++++++++++--
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py
index b6712d9602..b75f6066ad 100644
--- a/zerver/tests/test_retention.py
+++ b/zerver/tests/test_retention.py
@@ -4,10 +4,10 @@ from typing import Any, Dict, List, Optional
 
 from django.utils.timezone import now as timezone_now
 
-from zerver.lib.actions import internal_send_private_message
+from zerver.lib.actions import internal_send_private_message, do_add_submessage
 from zerver.lib.test_classes import ZulipTestCase
 from zerver.lib.upload import create_attachment
-from zerver.models import (Message, Realm, UserProfile, ArchivedUserMessage,
+from zerver.models import (Message, Realm, UserProfile, ArchivedUserMessage, SubMessage,
                            ArchivedMessage, Attachment, ArchivedAttachment, UserMessage,
                            get_realm, get_user_profile_by_email, get_system_bot)
 from zerver.lib.retention import (
@@ -301,6 +301,46 @@ class TestArchivingGeneral(RetentionTestingBase):
             sorted(msgs_ids.values())
         )
 
+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')
+
+        do_add_submessage(
+            realm=self.zulip_realm,
+            sender_id=cordelia.id,
+            message_id=expired_msg_ids[0],
+            msg_type='whatever',
+            content='{"name": "alice", "salary": 20}'
+        )
+        do_add_submessage(
+            realm=self.zulip_realm,
+            sender_id=hamlet.id,
+            message_id=expired_msg_ids[0],
+            msg_type='whatever',
+            content='{"name": "john", "salary": 30}'
+        )
+
+        do_add_submessage(
+            realm=self.zulip_realm,
+            sender_id=cordelia.id,
+            message_id=expired_msg_ids[1],
+            msg_type='whatever',
+            content='{"name": "jack", "salary": 10}'
+        )
+
+        submessage_ids = list(
+            SubMessage.objects.filter(message_id__in=expired_msg_ids).values_list('id', flat=True)
+        )
+
+        self.assertEqual(len(submessage_ids), 3)
+        self.assertEqual(SubMessage.objects.filter(id__in=submessage_ids).count(), 3)
+        archive_messages()
+        self.assertEqual(SubMessage.objects.filter(id__in=submessage_ids).count(), 0)
+
 class TestMoveMessageToArchive(ZulipTestCase):
     def setUp(self) -> None:
         self.sender = 'hamlet@zulip.com'