From 150f77aea2002aa4f9c174a8e73c9a3b83b71e34 Mon Sep 17 00:00:00 2001 From: Sahil Batra Date: Tue, 12 Jul 2022 17:12:29 +0530 Subject: [PATCH] bulk_create: Add users to system user groups in bulk_create_users. This commit modifies bulk_create_users to add the users to the respective system groups. And due to this change, now bots in development environment are also added to system groups. Tests are changed accordingly as more UserGroupMembeship objects are created. --- zerver/lib/bulk_create.py | 23 +++++++++++++++++++++ zerver/tests/test_realm.py | 4 ++-- zerver/tests/test_user_groups.py | 4 ++-- zilencer/management/commands/populate_db.py | 21 ------------------- 4 files changed, 27 insertions(+), 25 deletions(-) diff --git a/zerver/lib/bulk_create.py b/zerver/lib/bulk_create.py index b214146223..51bf290046 100644 --- a/zerver/lib/bulk_create.py +++ b/zerver/lib/bulk_create.py @@ -12,6 +12,8 @@ from zerver.models import ( Recipient, Stream, Subscription, + UserGroup, + UserGroupMembership, UserProfile, ) @@ -116,6 +118,27 @@ def bulk_create_users( Subscription.objects.bulk_create(subscriptions_to_create) + full_members_system_group = UserGroup.objects.get( + name="@role:fullmembers", realm=realm, is_system_group=True + ) + members_system_group = UserGroup.objects.get( + name="@role:members", realm=realm, is_system_group=True + ) + group_memberships_to_create: List[UserGroupMembership] = [] + for user_profile in profiles_to_create: + # All users are members since this function is only used to create bots + # and test and development environment users. + assert user_profile.role == UserProfile.ROLE_MEMBER + group_memberships_to_create.append( + UserGroupMembership(user_profile=user_profile, user_group=members_system_group) + ) + if not user_profile.is_provisional_member: + group_memberships_to_create.append( + UserGroupMembership(user_profile=user_profile, user_group=full_members_system_group) + ) + + UserGroupMembership.objects.bulk_create(group_memberships_to_create) + def bulk_set_users_or_streams_recipient_fields( model: Type[Model], diff --git a/zerver/tests/test_realm.py b/zerver/tests/test_realm.py index 2164cd0c5d..cc9126c97a 100644 --- a/zerver/tests/test_realm.py +++ b/zerver/tests/test_realm.py @@ -965,9 +965,9 @@ class RealmTest(ZulipTestCase): realm=realm, name="@role:fullmembers", is_system_group=True ) - self.assert_length(UserGroupMembership.objects.filter(user_group=members_system_group), 6) + self.assert_length(UserGroupMembership.objects.filter(user_group=members_system_group), 10) self.assert_length( - UserGroupMembership.objects.filter(user_group=full_members_system_group), 6 + UserGroupMembership.objects.filter(user_group=full_members_system_group), 10 ) self.assertEqual(realm.waiting_period_threshold, 0) diff --git a/zerver/tests/test_user_groups.py b/zerver/tests/test_user_groups.py index 58b11d216d..3678bdd569 100644 --- a/zerver/tests/test_user_groups.py +++ b/zerver/tests/test_user_groups.py @@ -309,11 +309,11 @@ class UserGroupAPITestCase(UserGroupTestCase): user_group = UserGroup.objects.get(name="support") # Test success self.assertEqual(UserGroup.objects.filter(realm=hamlet.realm).count(), 9) - self.assertEqual(UserGroupMembership.objects.count(), 19) + self.assertEqual(UserGroupMembership.objects.count(), 47) result = self.client_delete(f"/json/user_groups/{user_group.id}") self.assert_json_success(result) self.assertEqual(UserGroup.objects.filter(realm=hamlet.realm).count(), 8) - self.assertEqual(UserGroupMembership.objects.count(), 18) + self.assertEqual(UserGroupMembership.objects.count(), 46) # Test when invalid user group is supplied result = self.client_delete("/json/user_groups/1111") self.assert_json_error(result, "Invalid user group") diff --git a/zilencer/management/commands/populate_db.py b/zilencer/management/commands/populate_db.py index aec10bc517..35f8414292 100644 --- a/zilencer/management/commands/populate_db.py +++ b/zilencer/management/commands/populate_db.py @@ -56,8 +56,6 @@ from zerver.models import ( Service, Stream, Subscription, - UserGroup, - UserGroupMembership, UserMessage, UserPresence, UserProfile, @@ -500,25 +498,6 @@ class Command(BaseCommand): assign_time_zone_by_delivery_email("shiva@zulip.com", "Asia/Kolkata") # India assign_time_zone_by_delivery_email("cordelia@zulip.com", "UTC") - users = UserProfile.objects.filter(realm=zulip_realm) - # All users in development environment are full members initially because - # waiting period threshold is 0. Groups of Iago, Dedemona, Shiva and - # Polonius will be updated according to their role in do_change_user_role. - full_members_user_group = UserGroup.objects.get( - realm=zulip_realm, name="@role:fullmembers", is_system_group=True - ) - members_user_group = UserGroup.objects.get( - realm=zulip_realm, name="@role:members", is_system_group=True - ) - user_group_memberships = [] - for user_profile in list(users): - for group in [full_members_user_group, members_user_group]: - user_group_membership = UserGroupMembership( - user_group=group, user_profile=user_profile - ) - user_group_memberships.append(user_group_membership) - UserGroupMembership.objects.bulk_create(user_group_memberships) - iago = get_user_by_delivery_email("iago@zulip.com", zulip_realm) do_change_user_role(iago, UserProfile.ROLE_REALM_ADMINISTRATOR, acting_user=None) iago.is_staff = True