streams: Extract stream_traffic library.

This is a pure code move.

zerver/lib/actions.py

@@ -42,7 +42,7 @@ from psycopg2.sql import SQL
 from typing_extensions import TypedDict
 
 from analytics.lib.counts import COUNT_STATS, do_increment_logging_stat
-from analytics.models import RealmCount, StreamCount
+from analytics.models import RealmCount
 from confirmation import settings as confirmation_settings
 from confirmation.models import (
     Confirmation,
@@ -144,6 +144,7 @@ from zerver.lib.stream_subscription import (
     subscriber_ids_with_stream_history_access,
 )
 from zerver.lib.stream_topic import StreamTopicTarget
+from zerver.lib.stream_traffic import get_average_weekly_stream_traffic, get_streams_traffic
 from zerver.lib.streams import (
     access_stream_by_id,
     access_stream_for_send_message,
@@ -7163,51 +7164,6 @@ def do_delete_messages_by_sender(user: UserProfile) -> None:
         move_messages_to_archive(message_ids, chunk_size=retention.STREAM_MESSAGE_BATCH_SIZE)
 
 
-def get_streams_traffic(stream_ids: Set[int]) -> Dict[int, int]:
-    stat = COUNT_STATS["messages_in_stream:is_bot:day"]
-    traffic_from = timezone_now() - datetime.timedelta(days=28)
-
-    query = StreamCount.objects.filter(property=stat.property, end_time__gt=traffic_from)
-    query = query.filter(stream_id__in=stream_ids)
-
-    traffic_list = query.values("stream_id").annotate(value=Sum("value"))
-    traffic_dict = {}
-    for traffic in traffic_list:
-        traffic_dict[traffic["stream_id"]] = traffic["value"]
-
-    return traffic_dict
-
-
-def round_to_2_significant_digits(number: int) -> int:
-    return int(round(number, 2 - len(str(number))))
-
-
-STREAM_TRAFFIC_CALCULATION_MIN_AGE_DAYS = 7
-
-
-def get_average_weekly_stream_traffic(
-    stream_id: int, stream_date_created: datetime.datetime, recent_traffic: Dict[int, int]
-) -> Optional[int]:
-    try:
-        stream_traffic = recent_traffic[stream_id]
-    except KeyError:
-        stream_traffic = 0
-
-    stream_age = (timezone_now() - stream_date_created).days
-
-    if stream_age >= 28:
-        average_weekly_traffic = int(stream_traffic // 4)
-    elif stream_age >= STREAM_TRAFFIC_CALCULATION_MIN_AGE_DAYS:
-        average_weekly_traffic = int(stream_traffic * 7 // stream_age)
-    else:
-        return None
-
-    if average_weekly_traffic == 0 and stream_traffic > 0:
-        average_weekly_traffic = 1
-
-    return round_to_2_significant_digits(average_weekly_traffic)
-
-
 def get_web_public_subs(realm: Realm) -> SubscriptionInfo:
     color_idx = 0
 

zerver/lib/stream_traffic.py

@@ -0,0 +1,53 @@
+import datetime
+from typing import Dict, Optional, Set
+
+from django.db.models import Sum
+from django.utils.timezone import now as timezone_now
+
+from analytics.lib.counts import COUNT_STATS
+from analytics.models import StreamCount
+
+
+def get_streams_traffic(stream_ids: Set[int]) -> Dict[int, int]:
+    stat = COUNT_STATS["messages_in_stream:is_bot:day"]
+    traffic_from = timezone_now() - datetime.timedelta(days=28)
+
+    query = StreamCount.objects.filter(property=stat.property, end_time__gt=traffic_from)
+    query = query.filter(stream_id__in=stream_ids)
+
+    traffic_list = query.values("stream_id").annotate(value=Sum("value"))
+    traffic_dict = {}
+    for traffic in traffic_list:
+        traffic_dict[traffic["stream_id"]] = traffic["value"]
+
+    return traffic_dict
+
+
+def round_to_2_significant_digits(number: int) -> int:
+    return int(round(number, 2 - len(str(number))))
+
+
+STREAM_TRAFFIC_CALCULATION_MIN_AGE_DAYS = 7
+
+
+def get_average_weekly_stream_traffic(
+    stream_id: int, stream_date_created: datetime.datetime, recent_traffic: Dict[int, int]
+) -> Optional[int]:
+    try:
+        stream_traffic = recent_traffic[stream_id]
+    except KeyError:
+        stream_traffic = 0
+
+    stream_age = (timezone_now() - stream_date_created).days
+
+    if stream_age >= 28:
+        average_weekly_traffic = int(stream_traffic // 4)
+    elif stream_age >= STREAM_TRAFFIC_CALCULATION_MIN_AGE_DAYS:
+        average_weekly_traffic = int(stream_traffic * 7 // stream_age)
+    else:
+        return None
+
+    if average_weekly_traffic == 0 and stream_traffic > 0:
+        average_weekly_traffic = 1
+
+    return round_to_2_significant_digits(average_weekly_traffic)

zerver/tests/test_audit_log.py

@@ -37,9 +37,9 @@ from zerver.lib.actions import (
     do_set_realm_message_editing,
     do_set_realm_notifications_stream,
     do_set_realm_signup_notifications_stream,
-    get_streams_traffic,
 )
 from zerver.lib.message import get_last_message_id
+from zerver.lib.stream_traffic import get_streams_traffic
 from zerver.lib.streams import create_stream_if_needed
 from zerver.lib.test_classes import ZulipTestCase
 from zerver.models import (

zerver/tests/test_subs.py

@@ -35,12 +35,10 @@ from zerver.lib.actions import (
     ensure_stream,
     gather_subscriptions,
     gather_subscriptions_helper,
-    get_average_weekly_stream_traffic,
     get_default_streams_for_realm,
     get_topic_messages,
     lookup_default_stream_groups,
     pick_colors,
-    round_to_2_significant_digits,
     validate_user_access_to_subscribers_helper,
 )
 from zerver.lib.exceptions import JsonableError
@@ -51,6 +49,10 @@ from zerver.lib.stream_subscription import (
     num_subscribers_for_stream_id,
     subscriber_ids_with_stream_history_access,
 )
+from zerver.lib.stream_traffic import (
+    get_average_weekly_stream_traffic,
+    round_to_2_significant_digits,
+)
 from zerver.lib.streams import (
     StreamDict,
     access_stream_by_id,