From 926cd7c5bb22dd1670ee221de473c071e4aad1d8 Mon Sep 17 00:00:00 2001
From: John Baldwin <jlbaldwin@gmail.com>
Date: Thu, 17 Feb 2022 14:41:58 -0500
Subject: [PATCH] Enables alternate Celery daily task for Figures

This commit adds ENV_TOKENS['FIGURES']['DAILY_TASK'] override to define
which Celery task is called by CeleryBeat to run the daily data update
---
 figures/settings/lms_production.py |  6 +++-
 tests/test_settings.py             | 47 ++++++++++++++++++++++++++++--
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/figures/settings/lms_production.py b/figures/settings/lms_production.py
index b85002b0..fd7c3b5d 100644
--- a/figures/settings/lms_production.py
+++ b/figures/settings/lms_production.py
@@ -5,6 +5,8 @@
 import os
 from celery.schedules import crontab
 
+FIGURES_DEFAULT_DAILY_TASK = 'figures.tasks.populate_daily_metrics'
+
 
 class FiguresRouter(object):
 
@@ -64,8 +66,10 @@ def update_celerybeat_schedule(
     https://stackoverflow.com/questions/51631455/how-to-route-tasks-to-different-queues-with-celery-and-django
     """
     if figures_env_tokens.get('ENABLE_DAILY_METRICS_IMPORT', True):
+        figures_daily_task = figures_env_tokens.get('DAILY_TASK',
+                                                    FIGURES_DEFAULT_DAILY_TASK)
         celerybeat_schedule_settings['figures-populate-daily-metrics'] = {
-            'task': 'figures.tasks.populate_daily_metrics',
+            'task': figures_daily_task,
             'schedule': crontab(
                 hour=figures_env_tokens.get('DAILY_METRICS_IMPORT_HOUR', 2),
                 minute=figures_env_tokens.get('DAILY_METRICS_IMPORT_MINUTE', 0),
diff --git a/tests/test_settings.py b/tests/test_settings.py
index e9192377..8b106974 100644
--- a/tests/test_settings.py
+++ b/tests/test_settings.py
@@ -12,7 +12,10 @@
 
 
 from figures import helpers as figures_helpers
-from figures.settings.lms_production import plugin_settings
+from figures.settings.lms_production import (
+    FIGURES_DEFAULT_DAILY_TASK,
+    plugin_settings
+)
 
 
 @pytest.mark.parametrize('features, expected', [
@@ -90,6 +93,44 @@ def test_update_settings(self, figures_env_tokens, run_celery):
         assert settings.ENV_TOKENS['FIGURES'] == figures_env_tokens
 
 
+class TestDailyTaskFunction(object):
+    """Tests setting for Figures top level Celery daily job task
+
+    The top level Celery daily job task is what the scheduler calls to populate
+    Figures enrollment data snapshots and Figures daily aggregate metrics
+
+    We added this test case so that we can improve the daily Figures jobs, while
+    retaining the ability for deployments to run the previously implemented
+    daily jobs by default. This is a risk reduction strategy
+    """
+    ALTERNATE_TASK = 'figures.tasks.populate_daily_metrics_v2'
+
+    @pytest.fixture(autouse=True)
+    def setup(self, db):
+        self.settings = mock.Mock(
+            WEBPACK_LOADER={},
+            CELERYBEAT_SCHEDULE={},
+            FEATURES={},
+            ENV_TOKENS={},
+            CELERY_IMPORTS=[],
+        )
+
+    def test_uses_default_daily_task(self):
+        """Do the default settings define the default Celery task in CeleryBeat?
+        """
+        plugin_settings(self.settings)
+        task_found = self.settings.CELERYBEAT_SCHEDULE['figures-populate-daily-metrics']['task']
+        assert task_found == FIGURES_DEFAULT_DAILY_TASK
+
+    def test_uses_env_extra_daily_task(self):
+        """Does overriding the Celery task function setting in CeleryBeat?
+        """
+        self.settings.ENV_TOKENS['FIGURES'] = {'DAILY_TASK': self.ALTERNATE_TASK}
+        plugin_settings(self.settings)
+        task_found = self.settings.CELERYBEAT_SCHEDULE['figures-populate-daily-metrics']['task']
+        assert task_found == self.ALTERNATE_TASK
+
+
 class TestDailyMauPipelineSettings(object):
     """Tests MAU pipeline settings
 
@@ -115,14 +156,14 @@ def setup(self, db):
         )
 
     def test_daily_mau_pipeline_flag_enabled(self):
-        self.settings.ENV_TOKENS['FIGURES'] = { 'ENABLE_DAILY_MAU_IMPORT': True }
+        self.settings.ENV_TOKENS['FIGURES'] = {'ENABLE_DAILY_MAU_IMPORT': True}
         plugin_settings(self.settings)
         assert self.TASK_NAME in self.settings.CELERYBEAT_SCHEDULE
         assert set(['task', 'schedule', 'options']) == set(
             self.settings.CELERYBEAT_SCHEDULE[self.TASK_NAME].keys())
 
     def test_daily_mau_pipeline_flag_disabled(self):
-        self.settings.ENV_TOKENS['FIGURES'] = { 'ENABLE_DAILY_MAU_IMPORT': False }
+        self.settings.ENV_TOKENS['FIGURES'] = {'ENABLE_DAILY_MAU_IMPORT': False}
         plugin_settings(self.settings)
         assert self.TASK_NAME not in self.settings.CELERYBEAT_SCHEDULE