From 9a40319bfc4513db699f5983a9b233a7c5b97a27 Mon Sep 17 00:00:00 2001
From: Prakhar Pratyush <prakhar@zulip.com>
Date: Tue, 5 Nov 2024 20:10:42 +0530
Subject: [PATCH] management: Add option to reset all the active realms to a
 level.

This commit adds a --reset-level optional argument to
send_zulip_update_announcements management command to
reset all the active realms to a given level.
---
 .../commands/send_zulip_update_announcements.py    | 13 +++++++++++++
 zerver/tests/test_management_commands.py           | 14 ++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/zerver/management/commands/send_zulip_update_announcements.py b/zerver/management/commands/send_zulip_update_announcements.py
index 350883cfd8..fd3b9cdec4 100644
--- a/zerver/management/commands/send_zulip_update_announcements.py
+++ b/zerver/management/commands/send_zulip_update_announcements.py
@@ -1,10 +1,12 @@
 from argparse import ArgumentParser
 from typing import Any
 
+from django.conf import settings
 from typing_extensions import override
 
 from zerver.lib.management import ZulipBaseCommand, abort_unless_locked
 from zerver.lib.zulip_update_announcements import send_zulip_update_announcements
+from zerver.models import Realm
 
 
 class Command(ZulipBaseCommand):
@@ -17,8 +19,19 @@ class Command(ZulipBaseCommand):
             action="store_true",
             help="Immediately send updates if 'zulip_update_announcements_stream' is configured.",
         )
+        parser.add_argument(
+            "--reset-level",
+            type=int,
+            help="The level to reset all active realms to.",
+        )
 
     @override
     @abort_unless_locked
     def handle(self, *args: Any, **options: Any) -> None:
+        if options["reset_level"] is not None:
+            Realm.objects.filter(deactivated=False).exclude(
+                string_id=settings.SYSTEM_BOT_REALM
+            ).update(zulip_update_announcements_level=options["reset_level"])
+            return
+
         send_zulip_update_announcements(skip_delay=options["skip_delay"])
diff --git a/zerver/tests/test_management_commands.py b/zerver/tests/test_management_commands.py
index c91bbd4133..8e37df70df 100644
--- a/zerver/tests/test_management_commands.py
+++ b/zerver/tests/test_management_commands.py
@@ -588,3 +588,17 @@ class TestSendCustomEmail(ZulipTestCase):
                     call("  hamlet@zulip.com (zulip)"),
                 ],
             )
+
+
+class TestSendZulipUpdateAnnouncements(ZulipTestCase):
+    COMMAND_NAME = "send_zulip_update_announcements"
+
+    def test_reset_level(self) -> None:
+        realm = get_realm("zulip")
+        realm.zulip_update_announcements_level = 9
+        realm.save()
+
+        call_command(self.COMMAND_NAME, "--reset-level=5")
+
+        realm.refresh_from_db()
+        self.assertEqual(realm.zulip_update_announcements_level, 5)