From c5ac66b9c837644510034d83c42fe264f48a2db6 Mon Sep 17 00:00:00 2001 From: Mateusz Mandera Date: Fri, 24 May 2019 12:41:04 +0200 Subject: [PATCH] retention: Split archive_messages code into two functions. We split archive_messages code into two functions: moving to archive and cleanup. This allows cleaning up the tests - they can call these functions directly instead of copying several lines of archive_messages here and there in multiple tests. --- zerver/lib/retention.py | 9 +++++++-- zerver/tests/test_retention.py | 34 ++++++++++++---------------------- 2 files changed, 19 insertions(+), 24 deletions(-) diff --git a/zerver/lib/retention.py b/zerver/lib/retention.py index 492c4cfcb9..76394d773e 100644 --- a/zerver/lib/retention.py +++ b/zerver/lib/retention.py @@ -129,18 +129,23 @@ def clean_unused_messages() -> None: ) unused_messages.delete() - -def archive_messages() -> None: +def move_expired_to_archive() -> None: for realm in Realm.objects.filter(message_retention_days__isnull=False).order_by("id"): move_expired_messages_to_archive(realm) move_expired_user_messages_to_archive(realm) move_expired_attachments_to_archive(realm) move_expired_attachments_message_rows_to_archive(realm) + +def clean_expired() -> None: + for realm in Realm.objects.filter(message_retention_days__isnull=False).order_by("id"): delete_expired_user_messages(realm) delete_expired_messages(realm) delete_expired_attachments(realm) clean_unused_messages() +def archive_messages() -> None: + move_expired_to_archive() + clean_expired() def move_attachment_message_to_archive_by_message(message_ids: List[int]) -> None: # Move attachments messages relation table data to archive. diff --git a/zerver/tests/test_retention.py b/zerver/tests/test_retention.py index 414e7a1d59..41339f2489 100644 --- a/zerver/tests/test_retention.py +++ b/zerver/tests/test_retention.py @@ -12,11 +12,8 @@ from zerver.models import (Message, Realm, UserProfile, ArchivedUserMessage, get_user_profile_by_email, get_system_bot) from zerver.lib.retention import ( archive_messages, - clean_unused_messages, - delete_expired_messages, - delete_expired_user_messages, - move_expired_messages_to_archive, - move_expired_user_messages_to_archive, + clean_expired, + move_expired_to_archive, move_messages_to_archive ) @@ -123,10 +120,9 @@ class TestRetentionLib(ZulipTestCase): sent_message_id = self._send_cross_realm_message() all_user_messages_qty = UserMessage.objects.count() self._change_messages_pub_date([sent_message_id], timezone_now() - timedelta(days=period)) - realms = Realm.objects.filter(message_retention_days__isnull=False).order_by("id") - for realm_instance in realms: - move_expired_messages_to_archive(realm_instance) - move_expired_user_messages_to_archive(realm_instance) + + move_expired_to_archive() + user_messages_sent = UserMessage.objects.order_by('id').filter( message_id=sent_message_id) archived_messages = ArchivedMessage.objects.all() @@ -143,10 +139,8 @@ class TestRetentionLib(ZulipTestCase): [arc_user_msg.id for arc_user_msg in archived_user_messages], [user_msg.id for user_msg in user_messages_sent] ) - for realm_instance in realms: - delete_expired_user_messages(realm_instance) - delete_expired_messages(realm_instance) - clean_unused_messages() + + clean_expired() # Check messages and user messages after deleting expired messages # from the main tables. @@ -175,9 +169,8 @@ class TestRetentionLib(ZulipTestCase): } def test_no_expired_messages(self) -> None: - for realm_instance in Realm.objects.filter(message_retention_days__isnull=False): - move_expired_messages_to_archive(realm_instance) - move_expired_user_messages_to_archive(realm_instance) + move_expired_to_archive() + self.assertEqual(ArchivedUserMessage.objects.count(), 0) self.assertEqual(ArchivedMessage.objects.count(), 0) @@ -195,9 +188,8 @@ class TestRetentionLib(ZulipTestCase): zulip_msgs_ids, timezone_now() - timedelta(days=ZULIP_REALM_DAYS + 1)) - for realm_instance in Realm.objects.filter(message_retention_days__isnull=False): - move_expired_messages_to_archive(realm_instance) - move_expired_user_messages_to_archive(realm_instance) + move_expired_to_archive() + self.assertEqual(ArchivedMessage.objects.count(), len(expected_message_ids)) self.assertEqual( ArchivedUserMessage.objects.count(), @@ -213,9 +205,7 @@ class TestRetentionLib(ZulipTestCase): expected_mit_msgs = self._make_mit_messages( 5, timezone_now() - timedelta(days=MIT_REALM_DAYS + 1)) - for realm_instance in Realm.objects.filter(message_retention_days__isnull=False): - move_expired_messages_to_archive(realm_instance) - move_expired_user_messages_to_archive(realm_instance) + move_expired_to_archive() self.assertEqual(ArchivedMessage.objects.count(), 5) self.assertEqual(ArchivedUserMessage.objects.count(), 10)