mirror of
				https://github.com/zulip/zulip.git
				synced 2025-11-04 05:53:43 +00:00 
			
		
		
		
	We now update all test messages to have a pub_date of "now" in the setUp() function in TestRetentionLib. We've seen tests flake on query counts before this patch. It's not certain that the test flaked due to time-related glitches, but it seems the most plausible explanation.
		
			
				
	
	
		
			212 lines
		
	
	
		
			11 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
			
		
		
	
	
			212 lines
		
	
	
		
			11 KiB
		
	
	
	
		
			Python
		
	
	
	
	
	
# -*- coding: utf-8 -*-
 | 
						|
import types
 | 
						|
from datetime import datetime, timedelta
 | 
						|
 | 
						|
from django.utils.timezone import now as timezone_now
 | 
						|
from zerver.lib.test_classes import ZulipTestCase
 | 
						|
from zerver.lib.upload import create_attachment
 | 
						|
from zerver.models import Message, Realm, Recipient, UserProfile, UserMessage, ArchivedUserMessage, \
 | 
						|
    ArchivedMessage, Attachment, ArchivedAttachment
 | 
						|
from zerver.lib.retention import get_expired_messages, move_message_to_archive
 | 
						|
 | 
						|
from typing import Any, List
 | 
						|
 | 
						|
 | 
						|
class TestRetentionLib(ZulipTestCase):
 | 
						|
    """
 | 
						|
        Test receiving expired messages retention tool.
 | 
						|
    """
 | 
						|
 | 
						|
    def setUp(self) -> None:
 | 
						|
        super().setUp()
 | 
						|
        self.zulip_realm = self._set_realm_message_retention_value('zulip', 30)
 | 
						|
        self.mit_realm = self._set_realm_message_retention_value('zephyr', 100)
 | 
						|
        Message.objects.all().update(pub_date=timezone_now())
 | 
						|
 | 
						|
    @staticmethod
 | 
						|
    def _set_realm_message_retention_value(realm_str: str, retention_period: int) -> Realm:
 | 
						|
        realm = Realm.objects.get(string_id=realm_str)
 | 
						|
        realm.message_retention_days = retention_period
 | 
						|
        realm.save()
 | 
						|
        return realm
 | 
						|
 | 
						|
    @staticmethod
 | 
						|
    def _change_messages_pub_date(msgs_ids: List[int], pub_date: datetime) -> Any:
 | 
						|
        messages = Message.objects.filter(id__in=msgs_ids).order_by('id')
 | 
						|
        messages.update(pub_date=pub_date)
 | 
						|
        return messages
 | 
						|
 | 
						|
    def _make_mit_messages(self, message_quantity: int, pub_date: datetime) -> Any:
 | 
						|
        # send messages from mit.edu realm and change messages pub date
 | 
						|
        sender = self.mit_user('espuser')
 | 
						|
        recipient = self.mit_user('starnine')
 | 
						|
        msgs_ids = [self.send_personal_message(sender.email, recipient.email,
 | 
						|
                                               sender_realm='zephyr') for i in
 | 
						|
                    range(message_quantity)]
 | 
						|
        mit_messages = self._change_messages_pub_date(msgs_ids, pub_date)
 | 
						|
        return mit_messages
 | 
						|
 | 
						|
    def test_expired_messages_result_type(self) -> None:
 | 
						|
        # Check return type of get_expired_message method.
 | 
						|
        result = get_expired_messages()
 | 
						|
        self.assertIsInstance(result, types.GeneratorType)
 | 
						|
 | 
						|
    def test_no_expired_messages(self) -> None:
 | 
						|
        result = list(get_expired_messages())
 | 
						|
        self.assertFalse(result)
 | 
						|
 | 
						|
    def test_expired_messages_in_each_realm(self) -> None:
 | 
						|
        # Check result realm messages order and result content
 | 
						|
        # when all realm has expired messages.
 | 
						|
        expired_mit_messages = self._make_mit_messages(3, timezone_now() - timedelta(days=101))
 | 
						|
        self._make_mit_messages(4, timezone_now() - timedelta(days=50))
 | 
						|
        zulip_messages_ids = Message.objects.order_by('id').filter(
 | 
						|
            sender__realm=self.zulip_realm).values_list('id', flat=True)[3:10]
 | 
						|
        expired_zulip_messages = self._change_messages_pub_date(zulip_messages_ids,
 | 
						|
                                                                timezone_now() - timedelta(days=31))
 | 
						|
        # Iterate by result
 | 
						|
        expired_messages_result = [messages_list for messages_list in get_expired_messages()]
 | 
						|
        self.assertEqual(len(expired_messages_result), 2)
 | 
						|
        # Check mit.edu realm expired messages.
 | 
						|
        self.assertEqual(len(expired_messages_result[0]['expired_messages']), 3)
 | 
						|
        self.assertEqual(expired_messages_result[0]['realm_id'], self.mit_realm.id)
 | 
						|
        # Check zulip.com realm expired messages.
 | 
						|
        self.assertEqual(len(expired_messages_result[1]['expired_messages']), 7)
 | 
						|
        self.assertEqual(expired_messages_result[1]['realm_id'], self.zulip_realm.id)
 | 
						|
        # Compare expected messages ids with result messages ids.
 | 
						|
        self.assertEqual(
 | 
						|
            sorted([message.id for message in expired_mit_messages]),
 | 
						|
            [message.id for message in expired_messages_result[0]['expired_messages']]
 | 
						|
        )
 | 
						|
        self.assertEqual(
 | 
						|
            sorted([message.id for message in expired_zulip_messages]),
 | 
						|
            [message.id for message in expired_messages_result[1]['expired_messages']]
 | 
						|
        )
 | 
						|
 | 
						|
    def test_expired_messages_in_one_realm(self) -> None:
 | 
						|
        # Check realm with expired messages and messages
 | 
						|
        # with one day to expiration data.
 | 
						|
        expired_mit_messages = self._make_mit_messages(5, timezone_now() - timedelta(days=101))
 | 
						|
        actual_mit_messages = self._make_mit_messages(3, timezone_now() - timedelta(days=99))
 | 
						|
        expired_messages_result = list(get_expired_messages())
 | 
						|
        expired_mit_messages_ids = [message.id for message in expired_mit_messages]
 | 
						|
        expired_mit_messages_result_ids = [message.id for message in
 | 
						|
                                           expired_messages_result[0]['expired_messages']]
 | 
						|
        actual_mit_messages_ids = [message.id for message in actual_mit_messages]
 | 
						|
        self.assertEqual(len(expired_messages_result), 1)
 | 
						|
        self.assertEqual(len(expired_messages_result[0]['expired_messages']), 5)
 | 
						|
        self.assertEqual(expired_messages_result[0]['realm_id'], self.mit_realm.id)
 | 
						|
        # Compare expected messages ids with result messages ids.
 | 
						|
        self.assertEqual(
 | 
						|
            sorted(expired_mit_messages_ids),
 | 
						|
            expired_mit_messages_result_ids
 | 
						|
        )
 | 
						|
        # Check actual mit.edu messages are not contained in expired messages list
 | 
						|
        self.assertEqual(
 | 
						|
            set(actual_mit_messages_ids) - set(expired_mit_messages_ids),
 | 
						|
            set(actual_mit_messages_ids)
 | 
						|
        )
 | 
						|
 | 
						|
 | 
						|
class TestMoveMessageToArchive(ZulipTestCase):
 | 
						|
 | 
						|
    def setUp(self) -> None:
 | 
						|
        super().setUp()
 | 
						|
        self.sender = 'hamlet@zulip.com'
 | 
						|
        self.recipient = 'cordelia@zulip.com'
 | 
						|
 | 
						|
    def _create_attachments(self) -> None:
 | 
						|
        sample_size = 10
 | 
						|
        dummy_files = [
 | 
						|
            ('zulip.txt', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/zulip.txt', sample_size),
 | 
						|
            ('temp_file.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py', sample_size),
 | 
						|
            ('abc.py', '1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py', sample_size)
 | 
						|
        ]
 | 
						|
        user_profile = self.example_user('hamlet')
 | 
						|
        for file_name, path_id, size in dummy_files:
 | 
						|
            create_attachment(file_name, path_id, user_profile, size)
 | 
						|
 | 
						|
    def _check_messages_before_archiving(self, msg_id: int) -> List[int]:
 | 
						|
        user_messages_ids_before = list(UserMessage.objects.filter(
 | 
						|
            message_id=msg_id).order_by('id').values_list('id', flat=True))
 | 
						|
        self.assertEqual(ArchivedUserMessage.objects.count(), 0)
 | 
						|
        self.assertEqual(ArchivedMessage.objects.count(), 0)
 | 
						|
        return user_messages_ids_before
 | 
						|
 | 
						|
    def _check_messages_after_archiving(self, msg_id: int, user_msgs_ids_before: List[int]) -> None:
 | 
						|
        self.assertEqual(ArchivedMessage.objects.filter(id=msg_id).count(), 1)
 | 
						|
        self.assertEqual(Message.objects.filter(id=msg_id).count(), 0)
 | 
						|
        self.assertEqual(UserMessage.objects.filter(message_id=msg_id).count(), 0)
 | 
						|
        arc_user_messages_ids_after = list(ArchivedUserMessage.objects.filter(
 | 
						|
            message_id=msg_id).order_by('id').values_list('id', flat=True))
 | 
						|
        self.assertEqual(arc_user_messages_ids_after, user_msgs_ids_before)
 | 
						|
 | 
						|
    def test_personal_message_archiving(self) -> None:
 | 
						|
        msg_id = self.send_personal_message(self.sender, self.recipient)
 | 
						|
        user_messages_ids_before = self._check_messages_before_archiving(msg_id)
 | 
						|
        move_message_to_archive(message_id=msg_id)
 | 
						|
        self._check_messages_after_archiving(msg_id, user_messages_ids_before)
 | 
						|
 | 
						|
    def test_stream_message_archiving(self) -> None:
 | 
						|
        msg_id = self.send_stream_message(self.sender, "Verona")
 | 
						|
        user_messages_ids_before = self._check_messages_before_archiving(msg_id)
 | 
						|
        move_message_to_archive(message_id=msg_id)
 | 
						|
        self._check_messages_after_archiving(msg_id, user_messages_ids_before)
 | 
						|
 | 
						|
    def test_archiving_message_second_time(self) -> None:
 | 
						|
        msg_id = self.send_stream_message(self.sender, "Verona")
 | 
						|
        user_messages_ids_before = self._check_messages_before_archiving(msg_id)
 | 
						|
        move_message_to_archive(message_id=msg_id)
 | 
						|
        self._check_messages_after_archiving(msg_id, user_messages_ids_before)
 | 
						|
        with self.assertRaises(Message.DoesNotExist):
 | 
						|
            move_message_to_archive(message_id=msg_id)
 | 
						|
 | 
						|
    def test_archiving_message_with_attachment(self) -> None:
 | 
						|
        self._create_attachments()
 | 
						|
        body = """Some files here ...[zulip.txt](
 | 
						|
            http://localhost:9991/user_uploads/1/31/4CBjtTLYZhk66pZrF8hnYGwc/zulip.txt)
 | 
						|
            http://localhost:9991/user_uploads/1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py ....
 | 
						|
            Some more.... http://localhost:9991/user_uploads/1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py
 | 
						|
        """
 | 
						|
        msg_id = self.send_personal_message(self.sender, self.recipient, body)
 | 
						|
        user_messages_ids_before = self._check_messages_before_archiving(msg_id)
 | 
						|
        attachments_ids_before = list(Attachment.objects.filter(
 | 
						|
            messages__id=msg_id).order_by("id").values_list("id", flat=True))
 | 
						|
        self.assertEqual(ArchivedAttachment.objects.count(), 0)
 | 
						|
        move_message_to_archive(message_id=msg_id)
 | 
						|
        self._check_messages_after_archiving(msg_id, user_messages_ids_before)
 | 
						|
        self.assertEqual(Attachment.objects.count(), 0)
 | 
						|
        arc_attachments_ids_after = list(ArchivedAttachment.objects.filter(
 | 
						|
            messages__id=msg_id).order_by("id").values_list("id", flat=True))
 | 
						|
        self.assertEqual(attachments_ids_before, arc_attachments_ids_after)
 | 
						|
 | 
						|
    def test_archiving_message_with_shared_attachment(self) -> None:
 | 
						|
        # Check do not removing attachments which is used in other messages.
 | 
						|
        self._create_attachments()
 | 
						|
        body = """Some files here ...[zulip.txt](
 | 
						|
            http://localhost:9991/user_uploads/1/31/4CBjtTLYZhk66pZrF8hnYGwc/zulip.txt)
 | 
						|
            http://localhost:9991/user_uploads/1/31/4CBjtTLYZhk66pZrF8hnYGwc/temp_file.py ....
 | 
						|
            Some more.... http://localhost:9991/user_uploads/1/31/4CBjtTLYZhk66pZrF8hnYGwc/abc.py
 | 
						|
        """
 | 
						|
        msg_id = self.send_personal_message(self.sender, self.recipient, body)
 | 
						|
 | 
						|
        # Simulate a reply with the same contents.
 | 
						|
        msg_id_shared_attachments = self.send_personal_message(
 | 
						|
            from_email=self.recipient,
 | 
						|
            to_email=self.sender,
 | 
						|
            content=body,
 | 
						|
        )
 | 
						|
 | 
						|
        user_messages_ids_before = self._check_messages_before_archiving(msg_id)
 | 
						|
        attachments_ids_before = list(Attachment.objects.filter(
 | 
						|
            messages__id=msg_id).order_by("id").values_list("id", flat=True))
 | 
						|
        self.assertEqual(ArchivedAttachment.objects.count(), 0)
 | 
						|
        move_message_to_archive(message_id=msg_id)
 | 
						|
        self._check_messages_after_archiving(msg_id, user_messages_ids_before)
 | 
						|
        self.assertEqual(Attachment.objects.count(), 3)
 | 
						|
        arc_attachments_ids_after = list(ArchivedAttachment.objects.filter(
 | 
						|
            messages__id=msg_id).order_by("id").values_list("id", flat=True))
 | 
						|
        self.assertEqual(attachments_ids_before, arc_attachments_ids_after)
 | 
						|
        move_message_to_archive(message_id=msg_id_shared_attachments)
 | 
						|
        self.assertEqual(Attachment.objects.count(), 0)
 |