mirror of
				https://github.com/zulip/zulip.git
				synced 2025-11-04 05:53:43 +00:00 
			
		
		
		
	create_realm_internal_bots: Refactor to extract main op as a function.
We extract the entire operations of the management command to a function create_if_missing_realm_internal_bots in the zerver/lib/onboarding.py. The logic for determining if there are any realm internal bots which have not been created is extracted to a function missing_any_realm_internal_bots in actions.py.
This commit is contained in:
		@@ -4877,3 +4877,12 @@ def check_delete_user_group(user_group_id: int, user_profile: UserProfile) -> No
 | 
				
			|||||||
    user_group = access_user_group_by_id(user_group_id, user_profile)
 | 
					    user_group = access_user_group_by_id(user_group_id, user_profile)
 | 
				
			||||||
    user_group.delete()
 | 
					    user_group.delete()
 | 
				
			||||||
    do_send_delete_user_group_event(user_group_id, user_profile.realm.id)
 | 
					    do_send_delete_user_group_event(user_group_id, user_profile.realm.id)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					def missing_any_realm_internal_bots() -> bool:
 | 
				
			||||||
 | 
					    bot_emails = [bot['email_template'] % (settings.INTERNAL_BOT_DOMAIN,)
 | 
				
			||||||
 | 
					                  for bot in settings.REALM_INTERNAL_BOTS]
 | 
				
			||||||
 | 
					    bot_counts = dict(UserProfile.objects.filter(email__in=bot_emails)
 | 
				
			||||||
 | 
					                                         .values_list('email')
 | 
				
			||||||
 | 
					                                         .annotate(Count('id')))
 | 
				
			||||||
 | 
					    realm_count = Realm.objects.count()
 | 
				
			||||||
 | 
					    return any(bot_counts.get(email, 0) < realm_count for email in bot_emails)
 | 
				
			||||||
 
 | 
				
			|||||||
@@ -4,7 +4,7 @@ from django.conf import settings
 | 
				
			|||||||
from zerver.lib.actions import set_default_streams, bulk_add_subscriptions, \
 | 
					from zerver.lib.actions import set_default_streams, bulk_add_subscriptions, \
 | 
				
			||||||
    internal_prep_stream_message, internal_send_private_message, \
 | 
					    internal_prep_stream_message, internal_send_private_message, \
 | 
				
			||||||
    create_stream_if_needed, create_streams_if_needed, do_send_messages, \
 | 
					    create_stream_if_needed, create_streams_if_needed, do_send_messages, \
 | 
				
			||||||
    do_add_reaction_legacy, create_users
 | 
					    do_add_reaction_legacy, create_users, missing_any_realm_internal_bots
 | 
				
			||||||
from zerver.models import Realm, UserProfile, Message, Reaction, get_system_bot
 | 
					from zerver.models import Realm, UserProfile, Message, Reaction, get_system_bot
 | 
				
			||||||
 | 
					
 | 
				
			||||||
from typing import Any, Dict, List, Mapping
 | 
					from typing import Any, Dict, List, Mapping
 | 
				
			||||||
@@ -27,6 +27,15 @@ def setup_realm_internal_bots(realm: Realm) -> None:
 | 
				
			|||||||
        bot.bot_owner = bot
 | 
					        bot.bot_owner = bot
 | 
				
			||||||
        bot.save()
 | 
					        bot.save()
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					def create_if_missing_realm_internal_bots() -> None:
 | 
				
			||||||
 | 
					    """This checks if there is any realm internal bot missing.
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    If that is the case, it creates the missing realm internal bots.
 | 
				
			||||||
 | 
					    """
 | 
				
			||||||
 | 
					    if missing_any_realm_internal_bots():
 | 
				
			||||||
 | 
					        for realm in Realm.objects.all():
 | 
				
			||||||
 | 
					            setup_realm_internal_bots(realm)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
def send_initial_pms(user: UserProfile) -> None:
 | 
					def send_initial_pms(user: UserProfile) -> None:
 | 
				
			||||||
    organization_setup_text = ""
 | 
					    organization_setup_text = ""
 | 
				
			||||||
    if user.is_realm_admin:
 | 
					    if user.is_realm_admin:
 | 
				
			||||||
 
 | 
				
			|||||||
@@ -5,7 +5,7 @@ from django.conf import settings
 | 
				
			|||||||
from django.core.management.base import BaseCommand
 | 
					from django.core.management.base import BaseCommand
 | 
				
			||||||
from django.db.models import Count
 | 
					from django.db.models import Count
 | 
				
			||||||
 | 
					
 | 
				
			||||||
from zerver.lib.onboarding import setup_realm_internal_bots
 | 
					from zerver.lib.onboarding import create_if_missing_realm_internal_bots
 | 
				
			||||||
from zerver.models import Realm, UserProfile
 | 
					from zerver.models import Realm, UserProfile
 | 
				
			||||||
 | 
					
 | 
				
			||||||
class Command(BaseCommand):
 | 
					class Command(BaseCommand):
 | 
				
			||||||
@@ -16,19 +16,7 @@ These are normally created when the realm is, so this should be a no-op
 | 
				
			|||||||
except when upgrading to a version that adds a new realm internal bot.
 | 
					except when upgrading to a version that adds a new realm internal bot.
 | 
				
			||||||
"""
 | 
					"""
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    @staticmethod
 | 
					 | 
				
			||||||
    def missing_any_bots() -> bool:
 | 
					 | 
				
			||||||
        bot_emails = [bot['email_template'] % (settings.INTERNAL_BOT_DOMAIN,)
 | 
					 | 
				
			||||||
                      for bot in settings.REALM_INTERNAL_BOTS]
 | 
					 | 
				
			||||||
        bot_counts = dict(UserProfile.objects.filter(email__in=bot_emails)
 | 
					 | 
				
			||||||
                                             .values_list('email')
 | 
					 | 
				
			||||||
                                             .annotate(Count('id')))
 | 
					 | 
				
			||||||
        realm_count = Realm.objects.count()
 | 
					 | 
				
			||||||
        return any(bot_counts.get(email, 0) < realm_count for email in bot_emails)
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
    def handle(self, *args: Any, **options: Any) -> None:
 | 
					    def handle(self, *args: Any, **options: Any) -> None:
 | 
				
			||||||
        if self.missing_any_bots():
 | 
					        create_if_missing_realm_internal_bots()
 | 
				
			||||||
            for realm in Realm.objects.all():
 | 
					        # create_users is idempotent -- it's a no-op when a given email
 | 
				
			||||||
                setup_realm_internal_bots(realm)
 | 
					        # already has a user in a given realm.
 | 
				
			||||||
                # create_users is idempotent -- it's a no-op when a given email
 | 
					 | 
				
			||||||
                # already has a user in a given realm.
 | 
					 | 
				
			||||||
 
 | 
				
			|||||||
							
								
								
									
										25
									
								
								zerver/tests/test_onboarding.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										25
									
								
								zerver/tests/test_onboarding.py
									
									
									
									
									
										Normal file
									
								
							@@ -0,0 +1,25 @@
 | 
				
			|||||||
 | 
					# -*- coding: utf-8 -*-
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					from zerver.models import Realm, UserProfile
 | 
				
			||||||
 | 
					from zerver.lib.onboarding import create_if_missing_realm_internal_bots
 | 
				
			||||||
 | 
					from zerver.lib.test_classes import (
 | 
				
			||||||
 | 
					    ZulipTestCase,
 | 
				
			||||||
 | 
					)
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					class TestRealmInternalBotCreation(ZulipTestCase):
 | 
				
			||||||
 | 
					    def test_create_if_missing_realm_internal_bots(self) -> None:
 | 
				
			||||||
 | 
					        realm_internal_bots_dict = [{'var_name': 'TEST_BOT',
 | 
				
			||||||
 | 
					                                     'email_template': 'test-bot@%s',
 | 
				
			||||||
 | 
					                                     'name': 'Test Bot'}]
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					        def check_test_bot_exists() -> bool:
 | 
				
			||||||
 | 
					            all_realms_count = Realm.objects.count()
 | 
				
			||||||
 | 
					            all_test_bot_count = UserProfile.objects.filter(
 | 
				
			||||||
 | 
					                email='test-bot@zulip.com'
 | 
				
			||||||
 | 
					            ).count()
 | 
				
			||||||
 | 
					            return all_realms_count == all_test_bot_count
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					        self.assertFalse(check_test_bot_exists())
 | 
				
			||||||
 | 
					        with self.settings(REALM_INTERNAL_BOTS=realm_internal_bots_dict):
 | 
				
			||||||
 | 
					            create_if_missing_realm_internal_bots()
 | 
				
			||||||
 | 
					        self.assertTrue(check_test_bot_exists())
 | 
				
			||||||
		Reference in New Issue
	
	Block a user