From bb125602337d378782abc8190fb5131edf9d79b3 Mon Sep 17 00:00:00 2001
From: Calen Pennington <cale@edx.org>
Date: Thu, 23 Jul 2020 15:32:50 -0400
Subject: [PATCH] Make the ExperimentWaffleFlag respect course masquerading
 when checking if it's active for a specific enrollment

---
 lms/djangoapps/courseware/masquerade.py | 12 ++++++++++++
 lms/djangoapps/experiments/flags.py     | 23 ++++++++++++++++++-----
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/lms/djangoapps/courseware/masquerade.py b/lms/djangoapps/courseware/masquerade.py
index c679e937c4..35c4e5f74b 100644
--- a/lms/djangoapps/courseware/masquerade.py
+++ b/lms/djangoapps/courseware/masquerade.py
@@ -326,6 +326,18 @@ def is_masquerading_as_specific_student(user, course_key):
     return bool(course_masquerade and course_masquerade.user_name)
 
 
+def get_specific_masquerading_user(user, course_key):
+    """
+    Return the specific user that a staff member is masquerading as, or None if they aren't.
+    """
+    course_masquerade = get_course_masquerade(user, course_key)
+    is_specific_user = bool(course_masquerade and course_masquerade.user_name)
+    if is_specific_user:
+        return User.objects.get(username=course_masquerade.user_name)
+    else:
+        return None
+
+
 def get_masquerading_user_group(course_key, user, user_partition):
     """
     If the current user is masquerading as a generic learner in a specific group, return that group.
diff --git a/lms/djangoapps/experiments/flags.py b/lms/djangoapps/experiments/flags.py
index 2c6cd5ad16..2922a5303d 100644
--- a/lms/djangoapps/experiments/flags.py
+++ b/lms/djangoapps/experiments/flags.py
@@ -119,6 +119,12 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
         # Keep some imports in here, because this class is commonly used at a module level, and we want to avoid
         # circular imports for any models.
         from experiments.models import ExperimentKeyValue
+        from lms.djangoapps.courseware.access import has_access
+        from lms.djangoapps.courseware.masquerade import (
+            setup_masquerade,
+            is_masquerading,
+            get_specific_masquerading_user,
+        )
 
         request = get_current_request()
         if not request:
@@ -128,6 +134,13 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
             # We need username for stable bucketing and id for tracking, so just skip anonymous (not-logged-in) users
             return 0
 
+        user = get_specific_masquerading_user(request.user, course_key)
+        if user is None:
+            user = request.user
+            masquerading_as_specific_student = False
+        else:
+            masquerading_as_specific_student = True
+
         # Use course key in experiment name to separate caches and segment calls per-course-run
         experiment_name = self.namespaced_flag_name + ('.{}'.format(course_key) if course_key else '')
 
@@ -146,10 +159,10 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
             values = ExperimentKeyValue.objects.filter(experiment_id=self.experiment_id).values('key', 'value')
             values = {pair['key']: pair['value'] for pair in values}
 
-            if not self._is_enrollment_inside_date_bounds(values, request.user, course_key):
+            if not self._is_enrollment_inside_date_bounds(values, user, course_key):
                 return self._cache_bucket(experiment_name, 0)
 
-        bucket = stable_bucketing_hash_group(experiment_name, self.num_buckets, request.user.username)
+        bucket = stable_bucketing_hash_group(experiment_name, self.num_buckets, user.username)
 
         # Now check if the user is forced into a particular bucket, using our subordinate bucket flags
         for i, bucket_flag in enumerate(self.bucket_flags):
@@ -158,9 +171,9 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
                 break
 
         session_key = 'tracked.{}'.format(experiment_name)
-        if track and hasattr(request, 'session') and session_key not in request.session:
+        if track and hasattr(request, 'session') and session_key not in request.session and not masquerading_as_specific_student:
             segment.track(
-                user_id=request.user.id,
+                user_id=user.id,
                 event_name='edx.bi.experiment.user.bucketed',
                 properties={
                     'site': request.site.domain,
@@ -168,7 +181,7 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
                     'experiment': self.flag_name,
                     'course_id': str(course_key) if course_key else None,
                     'bucket': bucket,
-                    'is_staff': request.user.is_staff,
+                    'is_staff': user.is_staff,
                     'nonInteraction': 1,
                 }
             )