analytics: Refactor fixture creation to make it more general.

Also less verbose, in preparation for adding a bunch more fixtures.

analytics/lib/fixtures.py

@@ -66,11 +66,3 @@ def generate_time_series_data(days=100, business_hours_base=10, non_business_hou
         for i in range(1, length):
             values[i] = values[i-1] + values[i]
     return [max(v, 0) for v in values]
-
-def bulk_create_realmcount(property, subgroup, last_end_time, frequency, interval, values, realm):
-    # type: (str, str, datetime, str, str, List[int], Realm) -> None
-    end_times = time_range(last_end_time, last_end_time, frequency, len(values))
-    RealmCount.objects.bulk_create([
-        RealmCount(realm=realm, property=property, subgroup=subgroup, end_time=end_time,
-                   interval=interval, value=value)
-        for end_time, value in zip(end_times, values) if value != 0])

analytics/management/commands/populate_analytics_db.py

@@ -5,18 +5,24 @@ from argparse import ArgumentParser
 from django.core.management.base import BaseCommand
 from django.utils import timezone
 
-from analytics.models import InstallationCount, RealmCount, UserCount, StreamCount
+from analytics.models import BaseCount, InstallationCount, RealmCount, \
+    UserCount, StreamCount
 from analytics.lib.counts import COUNT_STATS, CountStat, do_drop_all_analytics_tables
-from analytics.lib.fixtures import generate_time_series_data, bulk_create_realmcount
+from analytics.lib.fixtures import generate_time_series_data
+from analytics.lib.time_utils import time_range
 from zerver.lib.timestamp import floor_to_day
 from zerver.models import Realm, UserProfile, Stream, Message
 
 from datetime import datetime, timedelta
-from typing import Any, Text
+
+from six.moves import zip
+from typing import Any, List, Optional, Text, Type, Union
 
 class Command(BaseCommand):
     help = """Populates analytics tables with randomly generated data."""
 
+    DAYS_OF_DATA = 100
+
     def create_user(self, email, full_name, is_staff, date_joined, realm):
         # type: (Text, Text, Text, bool, datetime, Realm) -> UserProfile
         return UserProfile.objects.create(
@@ -24,27 +30,43 @@ class Command(BaseCommand):
             realm=realm, short_name=full_name, pointer=-1, last_pointer_updater='none',
             api_key='42', date_joined=date_joined)
 
+    def generate_fixture_data(self, stat, business_hours_base, non_business_hours_base,
+                              growth, autocorrelation, spikiness, holiday_rate=0):
+        # type: (CountStat, float, float, float, float, float, float) -> List[int]
+        return generate_time_series_data(
+            days=self.DAYS_OF_DATA, business_hours_base=business_hours_base,
+            non_business_hours_base=non_business_hours_base, growth=growth,
+            autocorrelation=autocorrelation, spikiness=spikiness, holiday_rate=holiday_rate,
+            frequency=stat.frequency, is_gauge=(stat.interval == CountStat.GAUGE))
+
     def handle(self, *args, **options):
         # type: (*Any, **Any) -> None
         do_drop_all_analytics_tables()
         # I believe this also deletes any objects with this realm as a foreign key
         Realm.objects.filter(string_id='analytics').delete()
 
-        installation_time = timezone.now() - timedelta(days=100)
+        installation_time = timezone.now() - timedelta(days=self.DAYS_OF_DATA)
+        last_end_time = floor_to_day(timezone.now())
         realm = Realm.objects.create(
             string_id='analytics', name='Analytics', domain='analytics.ds',
             date_created=installation_time)
-        self.create_user('shylock@analytics.ds', 'Shylock', True, installation_time, realm)
+        shylock = self.create_user('shylock@analytics.ds', 'Shylock', True, installation_time, realm)
+
+        def insert_fixture_data(stat, fixture_data, table):
+            # type: (CountStat, Dict[Optional[str], List[int]], Type[BaseCount]) -> None
+            end_times = time_range(last_end_time, last_end_time, stat.frequency,
+                                   len(list(fixture_data.values())[0]))
+            if table == RealmCount:
+                id_args = {'realm': realm}
+            if table == UserCount:
+                id_args = {'realm': realm, 'user': shylock}
+            for subgroup, values in fixture_data.items():
+                table.objects.bulk_create([
+                    table(property=stat.property, subgroup=subgroup, end_time=end_time,
+                          interval=stat.interval, value=value, **id_args)
+                    for end_time, value in zip(end_times, values) if value != 0])
 
         stat = COUNT_STATS['active_users:is_bot']
-        if not RealmCount.objects.filter(property=stat.property).exists():
+        realm_data = {'false': self.generate_fixture_data(stat, .1, .03, 3, .5, 3),
-            last_end_time = floor_to_day(timezone.now())
+                      'true': self.generate_fixture_data(stat, .01, 0, 1, 0, 1)}
-            human_data = generate_time_series_data(days=100, business_hours_base=30,
+        insert_fixture_data(stat, realm_data, RealmCount)
-                                                   non_business_hours_base=10, growth=5, autocorrelation=.5,
-                                                   spikiness=3, frequency=CountStat.DAY)
-            bot_data = generate_time_series_data(days=100, business_hours_base=20,
-                                                 non_business_hours_base=20, growth=3, frequency=CountStat.DAY)
-            bulk_create_realmcount(stat.property, 'false', last_end_time,
-                                   stat.frequency, stat.interval, human_data, realm)
-            bulk_create_realmcount(stat.property, 'true', last_end_time,
-                                   stat.frequency, stat.interval, bot_data, realm)