From fb4b46c12c66709c46f17baf8a18b3a1ca70930a Mon Sep 17 00:00:00 2001
From: Awais Jibran <awaisdar001@gmail.com>
Date: Wed, 13 Nov 2019 16:02:50 +0500
Subject: [PATCH] Code Refactor for Course Grade Report

---
 .../instructor_task/config/waffle.py          | 15 ++++++++--
 .../instructor_task/tasks_helper/grades.py    | 28 +++++++++++--------
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/lms/djangoapps/instructor_task/config/waffle.py b/lms/djangoapps/instructor_task/config/waffle.py
index c5872c573d..d8d68f8015 100644
--- a/lms/djangoapps/instructor_task/config/waffle.py
+++ b/lms/djangoapps/instructor_task/config/waffle.py
@@ -4,10 +4,14 @@ waffle switches for the instructor_task app.
 """
 from __future__ import absolute_import
 
-from openedx.core.djangoapps.waffle_utils import CourseWaffleFlag, WaffleFlagNamespace
+from openedx.core.djangoapps.waffle_utils import CourseWaffleFlag, WaffleFlagNamespace, WaffleSwitchNamespace
 
+WAFFLE_NAMESPACE = u'instructor_task'
+INSTRUCTOR_TASK_WAFFLE_FLAG_NAMESPACE = WaffleFlagNamespace(name=WAFFLE_NAMESPACE)
+WAFFLE_SWITCHES = WaffleSwitchNamespace(name=WAFFLE_NAMESPACE)
 
-INSTRUCTOR_TASK_WAFFLE_FLAG_NAMESPACE = WaffleFlagNamespace(name=u'instructor_task')
+# Waffle switches
+OPTIMIZE_GET_LEARNERS_FOR_COURSE = u'optimize_get_learners_for_course'
 
 # Course-specific flags
 PROBLEM_GRADE_REPORT_VERIFIED_ONLY = u'problem_grade_report_verified_only'
@@ -48,3 +52,10 @@ def course_grade_report_verified_only(course_id):
     False otherwise.
     """
     return waffle_flags()[PROBLEM_GRADE_REPORT_VERIFIED_ONLY].is_enabled(course_id)
+
+
+def optimize_get_learners_switch_enabled():
+    """
+    Returns True if optimize get learner switch is enabled, otherwise False.
+    """
+    return WAFFLE_SWITCHES.is_enabled(OPTIMIZE_GET_LEARNERS_FOR_COURSE)
diff --git a/lms/djangoapps/instructor_task/tasks_helper/grades.py b/lms/djangoapps/instructor_task/tasks_helper/grades.py
index b25d2aac5b..d9e62cbce1 100644
--- a/lms/djangoapps/instructor_task/tasks_helper/grades.py
+++ b/lms/djangoapps/instructor_task/tasks_helper/grades.py
@@ -30,6 +30,7 @@ from lms.djangoapps.instructor_analytics.basic import list_problem_responses
 from lms.djangoapps.instructor_analytics.csvs import format_dictlist
 from lms.djangoapps.instructor_task.config.waffle import (
     course_grade_report_verified_only,
+    optimize_get_learners_switch_enabled,
     problem_grade_report_verified_only
 )
 from lms.djangoapps.teams.models import CourseTeamMembership
@@ -38,7 +39,6 @@ from opaque_keys.edx.keys import UsageKey
 from openedx.core.djangoapps.content.block_structure.api import get_course_in_cache
 from openedx.core.djangoapps.course_groups.cohorts import bulk_cache_cohorts, get_cohort, is_course_cohorted
 from openedx.core.djangoapps.user_api.course_tag.api import BulkCourseTags
-from openedx.core.djangoapps.waffle_utils import WaffleSwitchNamespace
 from student.models import CourseEnrollment
 from student.roles import BulkRoleCache
 from xmodule.modulestore.django import modulestore
@@ -48,10 +48,6 @@ from xmodule.split_test_module import get_split_user_partitions
 from .runner import TaskProgress
 from .utils import upload_csv_to_report_store
 
-WAFFLE_NAMESPACE = 'instructor_task'
-WAFFLE_SWITCHES = WaffleSwitchNamespace(name=WAFFLE_NAMESPACE)
-OPTIMIZE_GET_LEARNERS_FOR_COURSE = 'optimize_get_learners_for_course'
-
 TASK_LOG = logging.getLogger('edx.celery.task')
 
 ENROLLED_IN_COURSE = 'enrolled'
@@ -311,6 +307,20 @@ class CourseGradeReport(object):
             args = [iter(iterable)] * chunk_size
             return zip_longest(*args, fillvalue=fillvalue)
 
+        def get_enrolled_learners_for_course(course_id, verified_only=False):
+            """
+            Get enrolled learners in a course.
+
+            verified_only(bool): It indicates if we need only the verified
+            enrollments or all enrollments.
+            """
+            if optimize_get_learners_switch_enabled():
+                TASK_LOG.info(u'%s, Creating Course Grade with optimization', task_log_message)
+                return users_for_course_v2(course_id, verified_only=verified_only)
+
+            TASK_LOG.info(u'%s, Creating Course Grade without optimization', task_log_message)
+            return users_for_course(course_id, verified_only=verified_only)
+
         def users_for_course(course_id, verified_only=False):
             """
             Get all the enrolled users in a course.
@@ -356,13 +366,7 @@ class CourseGradeReport(object):
         course_id = context.course_id
         task_log_message = u'{}, Task type: {}'.format(context.task_info_string, context.action_name)
         verified_users_only = course_grade_report_verified_only(course_id)
-        if WAFFLE_SWITCHES.is_enabled(OPTIMIZE_GET_LEARNERS_FOR_COURSE):
-            TASK_LOG.info(u'%s, Creating Course Grade with optimization', task_log_message)
-            return users_for_course_v2(course_id, verified_only=verified_users_only)
-
-        TASK_LOG.info(u'%s, Creating Course Grade without optimization', task_log_message)
-        batch_users = users_for_course(course_id, verified_only=verified_users_only)
-        return batch_users
+        return get_enrolled_learners_for_course(course_id, verified_users_only)
 
     def _user_grades(self, course_grade, context):
         """