Allow specifying waffle switch on compute_grades

If --estimate_first_attempted is passed, it can force enabling or
disabling the ENABLE_FIRST_ATTEMPTED waffle switch.

TNL-6841

lms/djangoapps/grades/management/commands/compute_grades.py

@@ -17,6 +17,7 @@ from lms.djangoapps.grades.config.models import ComputeGradesSetting
 from student.models import CourseEnrollment
 from xmodule.modulestore.django import modulestore
 
+from ...config.waffle import waffle, ESTIMATE_FIRST_ATTEMPTED
 from ... import tasks
 
 
@@ -71,8 +72,15 @@ class Command(BaseCommand):
             default=0,
             type=int,
         )
+        parser.add_argument(
+            '--no_estimate_first_attempted',
+            help='Use score data to estimate first_attempted timestamp.',
+            action='store_false',
+            dest='estimate_first_attempted',
+        )
 
     def handle(self, *args, **options):
+
         self._set_log_level(options)
 
         for course_key in self._get_course_keys(options):
@@ -96,8 +104,9 @@ class Command(BaseCommand):
                 'course_key': six.text_type(course_key),
                 'offset': offset,
                 'batch_size': batch_size,
+                'estimate_first_attempted': options['estimate_first_attempted']
             }
-            result = tasks.compute_grades_for_course.apply_async(
+            result = tasks.compute_grades_for_course_v2.apply_async(
                 kwargs=kwargs,
                 options=task_options,
             )

lms/djangoapps/grades/management/commands/tests/test_compute_grades.py

@@ -17,6 +17,8 @@ from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
 from xmodule.modulestore.tests.factories import CourseFactory
 
 from lms.djangoapps.grades.config.models import ComputeGradesSetting
+from openedx.core.djangolib.waffle_utils import WaffleSwitchPlus
+from lms.djangoapps.grades.config.waffle import ESTIMATE_FIRST_ATTEMPTED
 from lms.djangoapps.grades.management.commands import compute_grades
 
 
@@ -72,35 +74,47 @@ class TestComputeGrades(SharedModuleStoreTestCase):
         with self.assertRaises(CommandError):
             self.command._get_course_keys({'from_settings': True})
 
-    @patch('lms.djangoapps.grades.tasks.compute_grades_for_course')
-    def test_tasks_fired(self, mock_task):
-        call_command(
+    @ddt.data(True, False)
+    @patch('lms.djangoapps.grades.tasks.compute_grades_for_course_v2')
+    def test_tasks_fired(self, estimate_first_attempted, mock_task):
+        command = [
             'compute_grades',
             '--routing_key=key',
             '--batch_size=2',
+        ]
+        courses = [
             '--courses',
             self.course_keys[0],
             self.course_keys[3],
             'd/n/e'  # No tasks created for nonexistent course, because it has no enrollments
-        )
+        ]
+        if not estimate_first_attempted:
+            command.append('--no_estimate_first_attempted')
+        call_command(*(command + courses))
+        _kwargs = lambda course_key, offset: {
+            'course_key': course_key,
+            'batch_size': 2,
+            'offset': offset,
+            'estimate_first_attempted': estimate_first_attempted
+        }
         self.assertEqual(
             mock_task.apply_async.call_args_list,
             [
                 ({
                     'options': {'routing_key': 'key'},
-                    'kwargs': {'course_key': self.course_keys[0], 'batch_size': 2, 'offset': 0}
+                    'kwargs': _kwargs(self.course_keys[0], 0)
                 },),
                 ({
                     'options': {'routing_key': 'key'},
-                    'kwargs': {'course_key': self.course_keys[0], 'batch_size': 2, 'offset': 2}
+                    'kwargs': _kwargs(self.course_keys[0], 2)
                 },),
                 ({
                     'options': {'routing_key': 'key'},
-                    'kwargs': {'course_key': self.course_keys[3], 'batch_size': 2, 'offset': 0}
+                    'kwargs': _kwargs(self.course_keys[3], 0)
                 },),
                 ({
                     'options': {'routing_key': 'key'},
-                    'kwargs': {'course_key': self.course_keys[3], 'batch_size': 2, 'offset': 2}
+                    'kwargs': _kwargs(self.course_keys[3], 2)
                 },),
             ],
         )

lms/djangoapps/grades/tasks.py

@@ -30,6 +30,7 @@ from track.event_transaction_utils import (
 from util.date_utils import from_timestamp
 from xmodule.modulestore.django import modulestore
 
+from .config.waffle import waffle, ESTIMATE_FIRST_ATTEMPTED
 from .constants import ScoreDatabaseTableEnum
 from .new.subsection_grade_factory import SubsectionGradeFactory
 from .new.course_grade_factory import CourseGradeFactory
@@ -61,7 +62,36 @@ class _BaseTask(PersistOnFailureTask, LoggedTask):  # pylint: disable=abstract-m
 
 
 @task(base=_BaseTask)
-def compute_grades_for_course(course_key, offset, batch_size):
+def compute_grades_for_course_v2(course_key, offset, batch_size, **kwargs):
+    """
+    Compute grades for a set of students in the specified course.
+
+    The set of students will be determined by the order of enrollment date, and
+    limited to at most <batch_size> students, starting from the specified
+    offset.
+
+    TODO: Roll this back into compute_grades_for_course once all workers have
+    the version with **kwargs.
+
+    Sets the ESTIMATE_FIRST_ATTEMPTED flag, then calls the original task as a
+    synchronous function.
+
+    estimate_first_attempted:
+        controls whether to unconditionally set the ESTIMATE_FIRST_ATTEMPTED
+        waffle switch.  If false or not provided, use the global value of
+        the ESTIMATE_FIRST_ATTEMPTED waffle switch.
+    """
+    course_key = kwargs.pop('course_key')
+    offset = kwargs.pop('offset')
+    batch_size = kwargs.pop('batch_size')
+    estimate_first_attempted = kwargs.pop('estimate_first_attempted', False)
+    if estimate_first_attempted:
+        waffle().override_for_request(ESTIMATE_FIRST_ATTEMPTED, True)
+    return compute_grades_for_course(course_key, offset, batch_size)
+
+
+@task(base=_BaseTask)
+def compute_grades_for_course(course_key, offset, batch_size, **kwargs):  # pylint: disable=unused-argument
     """
     Compute grades for a set of students in the specified course.