TNL-6739: adds in masquerading as enrollment track user

common/test/acceptance/pages/lms/staff_view.py

@@ -36,22 +36,22 @@ class StaffPreviewPage(PageObject):
     @property
     def staff_view_mode(self):
         """
-        Return the currently chosen view mode, e.g. "Staff", "Student" or a content group.
+        Return the currently chosen view mode, e.g. "Staff", "Learner" or a content group.
         """
         return self.q(css=self.VIEW_MODE_OPTIONS_CSS).filter(lambda el: el.is_selected()).first.text[0]
 
     def set_staff_view_mode(self, view_mode):
         """
-        Set the current view mode, e.g. "Staff", "Student" or a content group.
+        Set the current view mode, e.g. "Staff", "Learner" or a content group.
         """
         self.q(css=self.VIEW_MODE_OPTIONS_CSS).filter(lambda el: el.text.strip() == view_mode).first.click()
         self.wait_for_ajax()
 
     def set_staff_view_mode_specific_student(self, username_or_email):
         """
-        Set the current preview mode to "Specific Student" with the given username or email
+        Set the current preview mode to "Specific learner" with the given username or email
         """
-        required_mode = "Specific student"
+        required_mode = "Specific learner"
         if self.staff_view_mode != required_mode:
             self.q(css=self.VIEW_MODE_OPTIONS_CSS).filter(lambda el: el.text == required_mode).first.click()
         # Use a script here because .clear() + .send_keys() triggers unwanted behavior if a username is already set

common/test/acceptance/tests/lms/test_lms_cohorted_courseware_search.py

@@ -274,7 +274,7 @@ class CoursewareSearchCohortTest(ContainerBase):
         Test staff user can search just student public content if selected from preview menu.
         """
         self._auto_auth(self.staff_user["username"], self.staff_user["email"], False)
-        self._goto_staff_page().set_staff_view_mode('Student')
+        self._goto_staff_page().set_staff_view_mode('Learner')
         self.courseware_search_page.search_for_term(self.visible_to_all_html)
         assert self.visible_to_all_html in self.courseware_search_page.search_results.html[0]
         self.courseware_search_page.clear_search()
@@ -292,7 +292,7 @@ class CoursewareSearchCohortTest(ContainerBase):
         Test staff user can search cohort and public content if selected from preview menu.
         """
         self._auto_auth(self.staff_user["username"], self.staff_user["email"], False)
-        self._goto_staff_page().set_staff_view_mode('Student in ' + self.content_group_a)
+        self._goto_staff_page().set_staff_view_mode('Learner in ' + self.content_group_a)
         self.courseware_search_page.search_for_term(self.visible_to_all_html)
         assert self.visible_to_all_html in self.courseware_search_page.search_results.html[0]
         self.courseware_search_page.clear_search()

common/test/acceptance/tests/lms/test_lms_gating.py

@@ -203,7 +203,7 @@ class GatingTest(UniqueCourseTest):
         self.courseware_page.wait_for_page()
 
         self.course_home_page.visit()
-        self.course_home_page.preview.set_staff_view_mode('Student')
+        self.course_home_page.preview.set_staff_view_mode('Learner')
         self.assertEqual(self.course_home_page.outline.num_subsections, 1)
         self.course_home_page.outline.go_to_section('Test Section 1', 'Test Subsection 1')
         self.courseware_page.wait_for_page()

common/test/acceptance/tests/lms/test_lms_user_preview.py

 # -*- coding: utf-8 -*-
 """
-Tests the "preview" selector in the LMS that allows changing between Staff, Student, and Content Groups.
+Tests the "preview" selector in the LMS that allows changing between Staff, Learner, and Content Groups.
 """
 
 
@@ -13,7 +13,7 @@ from common.test.acceptance.pages.lms.instructor_dashboard import InstructorDash
 from common.test.acceptance.pages.lms.staff_view import StaffCoursewarePage
 from common.test.acceptance.fixtures.course import CourseFixture, XBlockFixtureDesc
 from bok_choy.promise import EmptyPromise
-from xmodule.partitions.partitions import Group
+from xmodule.partitions.partitions import (Group, ENROLLMENT_TRACK_PARTITION_ID, MINIMUM_STATIC_PARTITION_ID)
 from textwrap import dedent
 
 
@@ -98,8 +98,8 @@ class StaffViewToggleTest(CourseWithoutContentGroupsTest):
 
         course_page = self._goto_staff_page()
         self.assertTrue(course_page.has_tab('Instructor'))
-        course_page.set_staff_view_mode('Student')
-        self.assertEqual(course_page.staff_view_mode, 'Student')
+        course_page.set_staff_view_mode('Learner')
+        self.assertEqual(course_page.staff_view_mode, 'Learner')
         self.assertFalse(course_page.has_tab('Instructor'))
 
 
@@ -256,10 +256,13 @@ class CourseWithContentGroupsTest(StaffViewTest):
             "metadata": {
                 u"user_partitions": [
                     create_user_partition_json(
-                        0,
+                        MINIMUM_STATIC_PARTITION_ID,
                         'Configuration alpha,beta',
                         'Content Group Partition',
-                        [Group("0", 'alpha'), Group("1", 'beta')],
+                        [
+                            Group(MINIMUM_STATIC_PARTITION_ID + 1, 'alpha'),
+                            Group(MINIMUM_STATIC_PARTITION_ID + 2, 'beta')
+                        ],
                         scheme="cohort"
                     )
                 ],
@@ -285,6 +288,7 @@ class CourseWithContentGroupsTest(StaffViewTest):
 
         self.alpha_text = "VISIBLE TO ALPHA"
         self.beta_text = "VISIBLE TO BETA"
+        self.audit_text = "VISIBLE TO AUDIT"
         self.everyone_text = "VISIBLE TO EVERYONE"
 
         course_fixture.add_children(
@@ -292,12 +296,29 @@ class CourseWithContentGroupsTest(StaffViewTest):
                 XBlockFixtureDesc('sequential', 'Test Subsection').add_children(
                     XBlockFixtureDesc('vertical', 'Test Unit').add_children(
                         XBlockFixtureDesc(
-                            'problem', self.alpha_text, data=problem_data, metadata={"group_access": {0: [0]}}
+                            'problem',
+                            self.alpha_text,
+                            data=problem_data,
+                            metadata={"group_access": {MINIMUM_STATIC_PARTITION_ID: [MINIMUM_STATIC_PARTITION_ID + 1]}}
                         ),
                         XBlockFixtureDesc(
-                            'problem', self.beta_text, data=problem_data, metadata={"group_access": {0: [1]}}
+                            'problem',
+                            self.beta_text,
+                            data=problem_data,
+                            metadata={"group_access": {MINIMUM_STATIC_PARTITION_ID: [MINIMUM_STATIC_PARTITION_ID + 2]}}
                         ),
-                        XBlockFixtureDesc('problem', self.everyone_text, data=problem_data)
+                        XBlockFixtureDesc(
+                            'problem',
+                            self.audit_text,
+                            data=problem_data,
+                            # Below 1 is the hardcoded group ID for "Audit"
+                            metadata={"group_access": {ENROLLMENT_TRACK_PARTITION_ID: [1]}}
+                        ),
+                        XBlockFixtureDesc(
+                            'problem',
+                            self.everyone_text,
+                            data=problem_data
+                        )
                     )
                 )
             )
@@ -313,54 +334,72 @@ class CourseWithContentGroupsTest(StaffViewTest):
         Then I see all the problems, regardless of their group_access property
         """
         course_page = self._goto_staff_page()
-        verify_expected_problem_visibility(self, course_page, [self.alpha_text, self.beta_text, self.everyone_text])
+        verify_expected_problem_visibility(
+            self,
+            course_page,
+            [self.alpha_text, self.beta_text, self.audit_text, self.everyone_text]
+        )
 
     @attr(shard=3)
     def test_student_not_in_content_group(self):
         """
-        Scenario: When previewing as a student, only content visible to all is shown
+        Scenario: When previewing as a learner, only content visible to all is shown
         Given I have a course with a cohort user partition
         And problems that are associated with specific groups in the user partition
         When I view the courseware in the LMS with staff access
-        And I change to previewing as a Student
+        And I change to previewing as a Learner
         Then I see only problems visible to all users
         """
         course_page = self._goto_staff_page()
-        course_page.set_staff_view_mode('Student')
+        course_page.set_staff_view_mode('Learner')
         verify_expected_problem_visibility(self, course_page, [self.everyone_text])
 
     @attr(shard=3)
     def test_as_student_in_alpha(self):
         """
-        Scenario: When previewing as a student in group alpha, only content visible to alpha is shown
+        Scenario: When previewing as a learner in group alpha, only content visible to alpha is shown
         Given I have a course with a cohort user partition
         And problems that are associated with specific groups in the user partition
         When I view the courseware in the LMS with staff access
-        And I change to previewing as a Student in group alpha
+        And I change to previewing as a Learner in group alpha
         Then I see only problems visible to group alpha
         """
         course_page = self._goto_staff_page()
-        course_page.set_staff_view_mode('Student in alpha')
+        course_page.set_staff_view_mode('Learner in alpha')
         verify_expected_problem_visibility(self, course_page, [self.alpha_text, self.everyone_text])
 
     @attr(shard=3)
     def test_as_student_in_beta(self):
         """
-        Scenario: When previewing as a student in group beta, only content visible to beta is shown
+        Scenario: When previewing as a learner in group beta, only content visible to beta is shown
         Given I have a course with a cohort user partition
         And problems that are associated with specific groups in the user partition
         When I view the courseware in the LMS with staff access
-        And I change to previewing as a Student in group beta
+        And I change to previewing as a Learner in group beta
         Then I see only problems visible to group beta
         """
         course_page = self._goto_staff_page()
-        course_page.set_staff_view_mode('Student in beta')
+        course_page.set_staff_view_mode('Learner in beta')
         verify_expected_problem_visibility(self, course_page, [self.beta_text, self.everyone_text])
 
+    @attr(shard=3)
+    def test_as_student_in_audit(self):
+        """
+        Scenario: When previewing as a learner in the audit enrollment track, only content visible to audit is shown
+        Given I have a course with an enrollment_track user partition
+        And problems that are associated with specific groups in the user partition
+        When I view the courseware in the LMS with staff access
+        And I change to previewing as a Learner in audit enrollment track
+        Then I see only problems visible to audit enrollment track
+        """
+        course_page = self._goto_staff_page()
+        course_page.set_staff_view_mode('Learner in Audit')
+        verify_expected_problem_visibility(self, course_page, [self.audit_text, self.everyone_text])
+
     def create_cohorts_and_assign_students(self, student_a_username, student_b_username):
         """
         Adds 2 manual cohorts, linked to content groups, to the course.
-        Each cohort is assigned one student.
+        Each cohort is assigned one learner.
         """
         instructor_dashboard_page = InstructorDashboardPage(self.browser, self.course_id)
         instructor_dashboard_page.visit()
@@ -368,7 +407,7 @@ class CourseWithContentGroupsTest(StaffViewTest):
         cohort_management_page.is_cohorted = True
 
         def add_cohort_with_student(cohort_name, content_group, student):
-            """ Create cohort and assign student to it. """
+            """ Create cohort and assign learner to it. """
             cohort_management_page.add_cohort(cohort_name, content_group=content_group)
             cohort_management_page.add_students_to_selected_cohort([student])
         add_cohort_with_student("Cohort Alpha", "alpha", student_a_username)
@@ -383,14 +422,14 @@ class CourseWithContentGroupsTest(StaffViewTest):
         AutoAuthPage(self.browser, username=student_b_username, course_id=self.course_id, no_login=True).visit()
         self.create_cohorts_and_assign_students(student_a_username, student_b_username)
 
-        # Masquerade as student in alpha cohort:
+        # Masquerade as learner in alpha cohort:
         course_page = self._goto_staff_page()
         course_page.set_staff_view_mode_specific_student(student_a_username)
-        verify_expected_problem_visibility(self, course_page, [self.alpha_text, self.everyone_text])
+        verify_expected_problem_visibility(self, course_page, [self.alpha_text, self.audit_text, self.everyone_text])
 
-        # Masquerade as student in beta cohort:
+        # Masquerade as learner in beta cohort:
         course_page.set_staff_view_mode_specific_student(student_b_username)
-        verify_expected_problem_visibility(self, course_page, [self.beta_text, self.everyone_text])
+        verify_expected_problem_visibility(self, course_page, [self.beta_text, self.audit_text, self.everyone_text])
 
     @attr('a11y')
     def test_course_page(self):

common/test/acceptance/tests/studio/test_import_export.py

@@ -302,7 +302,7 @@ class TestEntranceExamCourseImport(ImportTestMixin, StudioCourseTest):
         self.landing_page.view_live()
         courseware = CoursewarePage(self.browser, self.course_id)
         courseware.wait_for_page()
-        StaffCoursewarePage(self.browser, self.course_id).set_staff_view_mode('Student')
+        StaffCoursewarePage(self.browser, self.course_id).set_staff_view_mode('Learner')
         self.assertEqual(courseware.num_sections, 1)
         self.assertIn(
             "To access course materials, you must score", courseware.entrance_exam_message_selector.text[0]

common/test/acceptance/tests/studio/test_studio_container.py

@@ -1076,14 +1076,14 @@ class UnitPublishingTest(ContainerBase):
         """
         Verifies no component is visible when viewing as a student.
         """
-        self._verify_and_return_staff_page().set_staff_view_mode('Student')
+        self._verify_and_return_staff_page().set_staff_view_mode('Learner')
         self.assertEqual(0, self.courseware.num_xblock_components)
 
     def _verify_student_view_visible(self, expected_components):
         """
         Verifies expected components are visible when viewing as a student.
         """
-        self._verify_and_return_staff_page().set_staff_view_mode('Student')
+        self._verify_and_return_staff_page().set_staff_view_mode('Learner')
         self._verify_components_visible(expected_components)
 
     def _verify_components_visible(self, expected_components):

common/test/acceptance/tests/studio/test_studio_outline.py

@@ -743,7 +743,7 @@ class StaffLockTest(CourseOutlineTest):
         course_home_page = CourseHomePage(self.browser, self.course_id)
         course_home_page.visit()
         self.assertEqual(course_home_page.outline.num_sections, 2)
-        course_home_page.preview.set_staff_view_mode('Student')
+        course_home_page.preview.set_staff_view_mode('Learner')
         self.assertEqual(course_home_page.outline.num_sections, 1)
 
     def test_locked_subsections_do_not_appear_in_lms(self):
@@ -764,7 +764,7 @@ class StaffLockTest(CourseOutlineTest):
         course_home_page = CourseHomePage(self.browser, self.course_id)
         course_home_page.visit()
         self.assertEqual(course_home_page.outline.num_subsections, 2)
-        course_home_page.preview.set_staff_view_mode('Student')
+        course_home_page.preview.set_staff_view_mode('Learner')
         self.assertEqual(course_home_page.outline.num_subsections, 1)
 
     def test_toggling_staff_lock_on_section_does_not_publish_draft_units(self):

lms/djangoapps/courseware/masquerade.py

@@ -17,6 +17,8 @@ from util.json_request import expect_json, JsonResponse
 from opaque_keys.edx.keys import CourseKey
 from xblock.fragment import Fragment
 from xblock.runtime import KeyValueStore
+from xmodule.partitions.partitions import NoSuchUserPartitionGroupError
+
 
 log = logging.getLogger(__name__)
 
@@ -170,16 +172,20 @@ def is_masquerading_as_specific_student(user, course_key):  # pylint: disable=in
     return bool(course_masquerade and course_masquerade.user_name)
 
 
-def get_masquerading_group_info(user, course_key):
+def get_masquerading_user_group(course_key, user, user_partition):
     """
-    If the user is masquerading as belonging to a group, then this method returns
-    two values: the id of the group, and the id of the user partition that the group
-    belongs to. If the user is not masquerading as a group, then None is returned.
+    If the current user is masquerading as a generic learner in a specific group, return that group.
+    If the user is not masquerading as a group, then None is returned.
     """
     course_masquerade = get_course_masquerade(user, course_key)
-    if not course_masquerade:
-        return None, None
-    return course_masquerade.group_id, course_masquerade.user_partition_id
+    if course_masquerade:
+        if course_masquerade.user_partition_id == user_partition.id and course_masquerade.group_id is not None:
+            try:
+                return user_partition.get_group(course_masquerade.group_id)
+            except NoSuchUserPartitionGroupError:
+                return None
+    # The user is masquerading as a generic student or not masquerading as a group return None
+    return None
 
 
 # Sentinel object to mark deleted objects in the session cache

lms/djangoapps/courseware/tests/test_masquerade.py

@@ -17,7 +17,7 @@ from courseware.masquerade import (
     MasqueradingKeyValueStore,
     handle_ajax,
     setup_masquerade,
-    get_masquerading_group_info
+    get_masquerading_user_group
 )
 from courseware.tests.factories import StaffFactory
 from courseware.tests.helpers import LoginEnrollmentTestCase, masquerade_as_group_member
@@ -435,22 +435,20 @@ class TestGetMasqueradingGroupId(StaffMasqueradeTestCase):
         modulestore().update_item(self.course, self.test_user.id)
 
     @patch.dict('django.conf.settings.FEATURES', {'DISABLE_START_DATES': False})
-    def test_group_masquerade(self):
+    def test_get_masquerade_group(self):
         """
-        Tests that a staff member can masquerade as being in a particular group.
+        Tests that a staff member can masquerade as being in a group in a user partition
         """
-        # Verify that there is no masquerading group initially
-        group_id, user_partition_id = get_masquerading_group_info(self.test_user, self.course.id)
-        self.assertIsNone(group_id)
-        self.assertIsNone(user_partition_id)
+        # Verify there is no masquerading group initially
+        group = get_masquerading_user_group(self.course.id, self.test_user, self.user_partition)
+        self.assertIsNone(group)
 
         # Install a masquerading group
         self.ensure_masquerade_as_group_member(0, 1)
 
         # Verify that the masquerading group is returned
-        group_id, user_partition_id = get_masquerading_group_info(self.test_user, self.course.id)
-        self.assertEqual(group_id, 1)
-        self.assertEqual(user_partition_id, 0)
+        group = get_masquerading_user_group(self.course.id, self.test_user, self.user_partition)
+        self.assertEqual(group.id, 1)
 
 
 class ReadOnlyKeyValueStore(DictKeyValueStore):

lms/static/lms/fixtures/preview/course_preview.html

@@ -6,10 +6,14 @@
                     <label for="action-preview-select" class="action-preview-label">View this course as:</label>
                     <select class="action-preview-select" id="action-preview-select" name="select">
                         <option value="staff" selected>Staff</option>
-                        <option value="student">Student</option>
-                        <option value="specific student">Specific student</option>
-                        <option value="group-a" data-group-id="group-a">Student in Group A</option>
-                        <option value="group-b" data-group-id="group-b">Student in Group B</option>
+                        <option value="student">Learner</option>
+                        <option value="specific student">Specific learner</option>
+                        <option value="group-a" data-group-id="group-a" data-partition-id="test_partition_a_id">
+                            Learner in Group A
+                        </option>
+                        <option value="group-b" data-group-id="group-b" data-partition-id="test_partition_b_id">
+                            Learner in Group B
+                        </option>
                     </select>
                     <div class="action-preview-username-container">
                         <label for="action-preview-username" class="action-preview-label">Username or email:</label>

lms/static/lms/js/preview/preview_factory.js

@@ -50,7 +50,7 @@
                 function masquerade(selectedOption) {
                     var data = {
                         role: selectedOption.val() === 'staff' ? 'staff' : 'student',
-                        user_partition_id: options.cohortedUserPartitionId,
+                        user_partition_id: selectedOption.data('partition-id'),
                         group_id: selectedOption.data('group-id'),
                         user_name: selectedOption.val() === 'specific student' ? $userNameElement.val() : null
                     };

lms/static/lms/js/spec/preview/preview_factory_spec.js

@@ -70,7 +70,6 @@ define(
                 var requests = AjaxHelpers.requests(this),
                     reloadSpy = spyOn(ViewUtils, 'reload');
                 showPreview({
-                    cohortedUserPartitionId: 'test_partition_id',
                     courseId: 'test_course'
                 });
                 previewActionSelect().find('option[value="group-b"]').prop('selected', 'selected').change();
@@ -79,7 +78,7 @@ define(
                     {
                         role: 'student',
                         user_name: null,
-                        user_partition_id: 'test_partition_id',
+                        user_partition_id: 'test_partition_b_id',
                         group_id: 'group-b'
                     }
                 );

lms/templates/preview_menu.html

@@ -5,9 +5,9 @@
 <%!
 from django.utils.translation import ugettext as _
 from django.conf import settings
-from openedx.core.djangoapps.course_groups.partition_scheme import get_cohorted_user_partition
 from openedx.core.djangolib.js_utils import dump_js_escaped_json
 from openedx.core.djangolib.markup import HTML, Text
+from xmodule.partitions.partitions_service import get_all_partitions_for_course
 %>
 
 <%
@@ -19,7 +19,7 @@ show_preview_menu = course and staff_access and supports_preview_menu
     def selected(is_selected):
       return "selected" if is_selected else ""
 
-    cohorted_user_partition = get_cohorted_user_partition(course)
+    course_partitions = get_all_partitions_for_course(course)
     masquerade_user_name = masquerade.user_name if masquerade else None
     masquerade_group_id = masquerade.group_id if masquerade else None
     staff_selected = selected(not masquerade or masquerade.role != "student")
@@ -34,13 +34,15 @@ show_preview_menu = course and staff_access and supports_preview_menu
                         <label for="action-preview-select" class="action-preview-label">${_("View this course as:")}</label>
                         <select class="action-preview-select" id="action-preview-select" name="select">
                             <option value="staff" ${staff_selected}>${_("Staff")}</option>
-                            <option value="student" ${student_selected}>${_("Student")}</option>
-                            <option value="specific student" ${specific_student_selected}>${_("Specific student")}</option>
-                            % if cohorted_user_partition:
-                                % for group in sorted(cohorted_user_partition.groups, key=lambda group: group.name):
-                                <option value="group.id" data-group-id="${group.id}" ${selected(masquerade_group_id == group.id)}>
-                                    ${_("Student in {content_group}").format(content_group=group.name)}
-                                </option>
+                            <option value="student" ${student_selected}>${_("Learner")}</option>
+                            <option value="specific student" ${specific_student_selected}>${_("Specific learner")}</option>
+                            % if course_partitions:
+                                % for course_partition in course_partitions:
+                                    % for group in sorted(course_partition.groups, key=lambda group: group.name):
+                                    <option value="group.id" data-group-id="${group.id}" data-partition-id="${course_partition.id}" ${selected(masquerade_group_id == group.id)}>
+                                        ${_("Learner in {content_group}").format(content_group=group.name)}
+                                    </option>
+                                    % endfor
                                 % endfor
                             % endif
                         </select>
@@ -71,7 +73,6 @@ show_preview_menu = course and staff_access and supports_preview_menu
         "courseId": course.id,
         "disableStudentAccess": disable_student_access if disable_student_access is not UNDEFINED else False,
         "specificStudentSelected": specific_student_selected,
-        "cohortedUserPartitionId": cohorted_user_partition.id if cohorted_user_partition else None,
         "masqueradeUsername" : masquerade_user_name if masquerade_user_name is not UNDEFINED else None,
     }
     %>

openedx/core/djangoapps/course_groups/partition_scheme.py

@@ -5,8 +5,8 @@ import logging
 
 from courseware.masquerade import (  # pylint: disable=import-error
     get_course_masquerade,
-    get_masquerading_group_info,
-    is_masquerading_as_specific_student,
+    get_masquerading_user_group,
+    is_masquerading_as_specific_student
 )
 from xmodule.partitions.partitions import NoSuchUserPartitionGroupError
 
@@ -45,14 +45,7 @@ class CohortPartitionScheme(object):
         # user is masquerading as a generic student in a specific group, then
         # return that group.
         if get_course_masquerade(user, course_key) and not is_masquerading_as_specific_student(user, course_key):
-            group_id, user_partition_id = get_masquerading_group_info(user, course_key)
-            if user_partition_id == user_partition.id and group_id is not None:
-                try:
-                    return user_partition.get_group(group_id)
-                except NoSuchUserPartitionGroupError:
-                    return None
-            # The user is masquerading as a generic student. We can't show any particular group.
-            return None
+            return get_masquerading_user_group(course_key, user, user_partition)
 
         cohort = get_cohort(user, course_key, use_cached=use_cached)
         if cohort is None:

openedx/core/djangoapps/verified_track_content/partition_scheme.py

@@ -5,8 +5,8 @@ from django.conf import settings
 
 from courseware.masquerade import (
     get_course_masquerade,
-    get_masquerading_group_info,
-    is_masquerading_as_specific_student,
+    get_masquerading_user_group,
+    is_masquerading_as_specific_student
 )
 from course_modes.models import CourseMode
 from student.models import CourseEnrollment
@@ -79,24 +79,13 @@ class EnrollmentTrackPartitionScheme(object):
         if is_course_using_cohort_instead(course_key):
             return None
 
-        # NOTE: masquerade code was copied from CohortPartitionScheme, and it may need
-        # some changes (or if not, code should be refactored out and shared).
-        # This work will be done in a future story TNL-6739.
-
         # First, check if we have to deal with masquerading.
         # If the current user is masquerading as a specific student, use the
         # same logic as normal to return that student's group. If the current
         # user is masquerading as a generic student in a specific group, then
         # return that group.
         if get_course_masquerade(user, course_key) and not is_masquerading_as_specific_student(user, course_key):
-            group_id, user_partition_id = get_masquerading_group_info(user, course_key)
-            if user_partition_id == user_partition.id and group_id is not None:
-                try:
-                    return user_partition.get_group(group_id)
-                except NoSuchUserPartitionGroupError:
-                    return None
-            # The user is masquerading as a generic student. We can't show any particular group.
-            return None
+            return get_masquerading_user_group(course_key, user, user_partition)
 
         mode_slug, is_active = CourseEnrollment.enrollment_mode_for_user(user, course_key)
         if mode_slug and is_active: