From 7137eba2229fa570ad580d55be10ed7ca835bc4a Mon Sep 17 00:00:00 2001 From: Sahil Batra Date: Wed, 16 Aug 2023 17:25:10 +0530 Subject: [PATCH] streams: Don't compute traffic data for sub objects in zephyr realm. We set stream_weekly_traffic field to "null" for Subscription objects in zephyr mirror realm as we do not need stream traffic data in zephyr mirror realm. This makes the subscription data consistent with steams data. This commit also udpates test to check never_subscribed data for zephyr mirror realm. --- zerver/actions/streams.py | 2 +- zerver/lib/subscription_info.py | 26 ++++++++++++++++---------- zerver/tests/test_subs.py | 17 +++++++++++++++-- 3 files changed, 32 insertions(+), 13 deletions(-) diff --git a/zerver/actions/streams.py b/zerver/actions/streams.py index 292c94df70..e205e2b493 100644 --- a/zerver/actions/streams.py +++ b/zerver/actions/streams.py @@ -328,7 +328,7 @@ def send_subscription_add_events( info_by_user[sub_info.user.id].append(sub_info) stream_ids = {sub_info.stream.id for sub_info in sub_info_list} - recent_traffic = get_streams_traffic(stream_ids=stream_ids) + recent_traffic = get_streams_traffic(stream_ids=stream_ids, realm=realm) # We generally only have a few streams, so we compute stream # data in its own loop. diff --git a/zerver/lib/subscription_info.py b/zerver/lib/subscription_info.py index df1346eecf..971e458073 100644 --- a/zerver/lib/subscription_info.py +++ b/zerver/lib/subscription_info.py @@ -109,7 +109,7 @@ def build_stream_dict_for_sub( user: UserProfile, sub_dict: RawSubscriptionDict, raw_stream_dict: RawStreamDict, - recent_traffic: Dict[int, int], + recent_traffic: Optional[Dict[int, int]], ) -> SubscriptionStreamDict: # Handle Stream.API_FIELDS can_remove_subscribers_group_id = raw_stream_dict["can_remove_subscribers_group_id"] @@ -145,9 +145,12 @@ def build_stream_dict_for_sub( is_announcement_only = raw_stream_dict["stream_post_policy"] == Stream.STREAM_POST_POLICY_ADMINS # Add a few computed fields not directly from the data models. - stream_weekly_traffic = get_average_weekly_stream_traffic( - raw_stream_dict["id"], raw_stream_dict["date_created"], recent_traffic - ) + if recent_traffic is not None: + stream_weekly_traffic = get_average_weekly_stream_traffic( + raw_stream_dict["id"], raw_stream_dict["date_created"], recent_traffic + ) + else: + stream_weekly_traffic = None email_address = encode_email_address_helper( raw_stream_dict["name"], raw_stream_dict["email_token"], show_sender=True @@ -184,7 +187,7 @@ def build_stream_dict_for_sub( def build_stream_dict_for_never_sub( raw_stream_dict: RawStreamDict, - recent_traffic: Dict[int, int], + recent_traffic: Optional[Dict[int, int]], ) -> NeverSubscribedStreamDict: can_remove_subscribers_group_id = raw_stream_dict["can_remove_subscribers_group_id"] date_created = datetime_to_timestamp(raw_stream_dict["date_created"]) @@ -198,9 +201,13 @@ def build_stream_dict_for_never_sub( rendered_description = raw_stream_dict["rendered_description"] stream_id = raw_stream_dict["id"] stream_post_policy = raw_stream_dict["stream_post_policy"] - stream_weekly_traffic = get_average_weekly_stream_traffic( - raw_stream_dict["id"], raw_stream_dict["date_created"], recent_traffic - ) + + if recent_traffic is not None: + stream_weekly_traffic = get_average_weekly_stream_traffic( + raw_stream_dict["id"], raw_stream_dict["date_created"], recent_traffic + ) + else: + stream_weekly_traffic = None # Backwards-compatibility addition of removed field. is_announcement_only = raw_stream_dict["stream_post_policy"] == Stream.STREAM_POST_POLICY_ADMINS @@ -430,8 +437,7 @@ def gather_subscriptions_helper( return recip_id_to_stream_id[sub_dict["recipient_id"]] traffic_stream_ids = {get_stream_id(sub_dict) for sub_dict in sub_dicts} - recent_traffic = get_streams_traffic(stream_ids=traffic_stream_ids) - assert recent_traffic is not None + recent_traffic = get_streams_traffic(stream_ids=traffic_stream_ids, realm=realm) # Okay, now we finally get to populating our main results, which # will be these three lists. diff --git a/zerver/tests/test_subs.py b/zerver/tests/test_subs.py index b57121671b..28274c1b71 100644 --- a/zerver/tests/test_subs.py +++ b/zerver/tests/test_subs.py @@ -4967,7 +4967,7 @@ class SubscriptionAPITest(ZulipTestCase): # realm. This does generate stream creation events from # send_stream_creation_events_for_previously_inaccessible_streams. with self.capture_send_event_calls(expected_num_events=num_streams + 1) as events: - with self.assert_database_query_count(num_streams + 12): + with self.assert_database_query_count(num_streams + 11): self.common_subscribe_to_streams( mit_user, stream_names, @@ -6249,7 +6249,7 @@ class GetSubscribersTest(ZulipTestCase): subdomain="zephyr", ) - with self.assert_database_query_count(4): + with self.assert_database_query_count(3): subscribed_streams, _ = gather_subscriptions(mit_user_profile, include_subscribers=True) self.assertGreaterEqual(len(subscribed_streams), 2) @@ -6260,6 +6260,19 @@ class GetSubscribersTest(ZulipTestCase): self.assert_length(sub["subscribers"], len(users_to_subscribe)) else: self.assert_length(sub["subscribers"], 0) + self.assertIsNone(sub["stream_weekly_traffic"]) + + # Create a web-public stream to test never_subscried data. + self.make_stream("mit_stream_2", realm=mit_user_profile.realm, is_web_public=True) + self.make_stream("mit_stream_3", realm=mit_user_profile.realm) + + sub_info = gather_subscriptions_helper(mit_user_profile, include_subscribers=True) + never_subscribed_streams = sub_info.never_subscribed + # Users in zephyr mirror realm can only access web-public never subscribed streams. + self.assert_length(never_subscribed_streams, 1) + self.assertEqual(never_subscribed_streams[0]["name"], "mit_stream_2") + self.assertTrue(never_subscribed_streams[0]["is_web_public"]) + self.assertIsNone(never_subscribed_streams[0]["stream_weekly_traffic"]) def test_nonsubscriber(self) -> None: """