From 322f6d1be2cbb96b57796b413fd388fb074c40a2 Mon Sep 17 00:00:00 2001
From: Daniel Friedman <dfriedman58@gmail.com>
Date: Fri, 30 Jan 2015 13:57:16 -0500
Subject: [PATCH] Ignore staff member's cohort when masquerade is on

TNL-1269
---
 .../course_groups/partition_scheme.py         | 18 ++++----
 .../tests/test_partition_scheme.py            | 41 ++++++++++++++++---
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/openedx/core/djangoapps/course_groups/partition_scheme.py b/openedx/core/djangoapps/course_groups/partition_scheme.py
index bb1530ac38..7249699a3a 100644
--- a/openedx/core/djangoapps/course_groups/partition_scheme.py
+++ b/openedx/core/djangoapps/course_groups/partition_scheme.py
@@ -36,15 +36,17 @@ class CohortPartitionScheme(object):
         If the user has no cohort mapping, or there is no (valid) cohort ->
         partition group mapping found, the function returns None.
         """
-        # If the current user is masquerading as being in a group belonging to the
-        # specified user partition then return the masquerading group.
+        # If the current user is masquerading as being in a group
+        # belonging to the specified user partition, return the
+        # masquerading group or None if the group can't be found.
         group_id, user_partition_id = get_masquerading_group_info(user, course_key)
-        if group_id is not None and user_partition_id == user_partition.id:
-            try:
-                return user_partition.get_group(group_id)
-            except NoSuchUserPartitionGroupError:
-                # If the group no longer exists then the masquerade is not in effect
-                pass
+        if user_partition_id == user_partition.id:
+            if group_id is not None:
+                try:
+                    return user_partition.get_group(group_id)
+                except NoSuchUserPartitionGroupError:
+                    return None
+            return None
 
         cohort = get_cohort(user, course_key)
         if cohort is None:
diff --git a/openedx/core/djangoapps/course_groups/tests/test_partition_scheme.py b/openedx/core/djangoapps/course_groups/tests/test_partition_scheme.py
index 574ec2bf17..1bce300b06 100644
--- a/openedx/core/djangoapps/course_groups/tests/test_partition_scheme.py
+++ b/openedx/core/djangoapps/course_groups/tests/test_partition_scheme.py
@@ -348,10 +348,9 @@ class TestMasqueradedGroup(StaffMasqueradeTestCase):
         # Send the request to set the masquerade
         request_json = {
             "role": "student",
+            "user_partition_id": self.user_partition.id,
+            "group_id": group.id if group is not None else None
         }
-        if group and self.user_partition:
-            request_json['user_partition_id'] = self.user_partition.id
-            request_json['group_id'] = group.id
         request = self._create_mock_json_request(
             self.test_user,
             body=json.dumps(request_json),
@@ -367,12 +366,42 @@ class TestMasqueradedGroup(StaffMasqueradeTestCase):
             group
         )
 
+    def _verify_masquerade_for_all_groups(self):
+        """
+        Verify that the staff user can masquerade as being in all groups
+        as well as no group.
+        """
+        self._verify_masquerade_for_group(self.user_partition.groups[0])
+        self._verify_masquerade_for_group(self.user_partition.groups[1])
+        self._verify_masquerade_for_group(None)
+
     @skipUnless(settings.ROOT_URLCONF == 'lms.urls', 'Test only valid in LMS')
     @patch.dict('django.conf.settings.FEATURES', {'DISABLE_START_DATES': False})
     def test_group_masquerade(self):
         """
         Tests that a staff member can masquerade as being in a particular group.
         """
-        self._verify_masquerade_for_group(self.user_partition.groups[0])
-        self._verify_masquerade_for_group(self.user_partition.groups[1])
-        self._verify_masquerade_for_group(None)
+        self._verify_masquerade_for_all_groups()
+
+    @skipUnless(settings.ROOT_URLCONF == 'lms.urls', 'Test only valid in LMS')
+    @patch.dict('django.conf.settings.FEATURES', {'DISABLE_START_DATES': False})
+    def test_group_masquerade_with_cohort(self):
+        """
+        Tests that a staff member can masquerade as being in a particular group
+        when that staff member also belongs to a cohort with a corresponding
+        group.
+        """
+        self.course.cohort_config = {'cohorted': True}
+        self.update_course(self.course, self.test_user.id)
+        cohort = CohortFactory.create(course_id=self.course.id, users=[self.test_user])
+        CourseUserGroupPartitionGroup(
+            course_user_group=cohort,
+            partition_id=self.user_partition.id,
+            group_id=self.user_partition.groups[0].id
+        ).save()
+
+        # When the staff user is masquerading as being in a None group
+        # (within an existent UserPartition), we should treat that as
+        # an explicit None, not defaulting to the user's cohort's
+        # partition group.
+        self._verify_masquerade_for_all_groups()