From 08ed3b547f8b042a3f42968066b4b307c667994d Mon Sep 17 00:00:00 2001
From: Eric Fischer <efischer@edx.org>
Date: Fri, 23 Oct 2015 15:16:30 -0400
Subject: [PATCH] CohortMemberships Unit Test Updates

Updating previously-existing unit tests to function with the new
CohortMembership model.
---
 .../instructor_task/tasks_helper.py           |  3 +--
 .../tests/test_tasks_helper.py                | 20 ++++++++++++-------
 .../djangoapps/course_groups/tests/helpers.py | 11 +++++++++-
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/lms/djangoapps/instructor_task/tasks_helper.py b/lms/djangoapps/instructor_task/tasks_helper.py
index a35199fd8c..8673b43ac0 100644
--- a/lms/djangoapps/instructor_task/tasks_helper.py
+++ b/lms/djangoapps/instructor_task/tasks_helper.py
@@ -1486,8 +1486,7 @@ def cohort_students_and_upload(_xmodule_instance_args, _entry_id, course_id, tas
                 continue
 
             try:
-                with transaction.commit_on_success():
-                    add_user_to_cohort(cohorts_status[cohort_name]['cohort'], username_or_email)
+                add_user_to_cohort(cohorts_status[cohort_name]['cohort'], username_or_email)
                 cohorts_status[cohort_name]['Students Added'] += 1
                 task_progress.succeeded += 1
             except User.DoesNotExist:
diff --git a/lms/djangoapps/instructor_task/tests/test_tasks_helper.py b/lms/djangoapps/instructor_task/tests/test_tasks_helper.py
index cc79806ad5..1645e8f315 100644
--- a/lms/djangoapps/instructor_task/tests/test_tasks_helper.py
+++ b/lms/djangoapps/instructor_task/tests/test_tasks_helper.py
@@ -20,7 +20,7 @@ from certificates.tests.factories import GeneratedCertificateFactory, Certificat
 from course_modes.models import CourseMode
 from courseware.tests.factories import InstructorFactory
 from instructor_task.tests.test_base import InstructorTaskCourseTestCase, TestReportMixin, InstructorTaskModuleTestCase
-from openedx.core.djangoapps.course_groups.models import CourseUserGroupPartitionGroup
+from openedx.core.djangoapps.course_groups.models import CourseUserGroupPartitionGroup, CohortMembership
 from openedx.core.djangoapps.course_groups.tests.helpers import CohortFactory
 import openedx.core.djangoapps.user_api.course_tag.api as course_tag_api
 from openedx.core.djangoapps.user_api.partition_schemes import RandomUserPartitionScheme
@@ -144,8 +144,10 @@ class TestInstructorGradeReport(InstructorGradeReportTestCase):
         magneto = u'MàgnëtÖ'
         cohort1 = CohortFactory(course_id=course.id, name=professor_x)
         cohort2 = CohortFactory(course_id=course.id, name=magneto)
-        cohort1.users.add(user1)
-        cohort2.users.add(user2)
+        membership1 = CohortMembership(course_user_group=cohort1, user=user1)
+        membership1.save()
+        membership2 = CohortMembership(course_user_group=cohort2, user=user2)
+        membership2.save()
 
         self._verify_cell_data_for_user(user1.username, course.id, 'Cohort Name', professor_x)
         self._verify_cell_data_for_user(user2.username, course.id, 'Cohort Name', magneto)
@@ -1399,8 +1401,10 @@ class TestCohortStudents(TestReportMixin, InstructorTaskCourseTestCase):
         )
 
     def test_move_users_to_new_cohort(self):
-        self.cohort_1.users.add(self.student_1)
-        self.cohort_2.users.add(self.student_2)
+        membership1 = CohortMembership(course_user_group=self.cohort_1, user=self.student_1)
+        membership1.save()
+        membership2 = CohortMembership(course_user_group=self.cohort_2, user=self.student_2)
+        membership2.save()
 
         result = self._cohort_students_and_upload(
             u'username,email,cohort\n'
@@ -1417,8 +1421,10 @@ class TestCohortStudents(TestReportMixin, InstructorTaskCourseTestCase):
         )
 
     def test_move_users_to_same_cohort(self):
-        self.cohort_1.users.add(self.student_1)
-        self.cohort_2.users.add(self.student_2)
+        membership1 = CohortMembership(course_user_group=self.cohort_1, user=self.student_1)
+        membership1.save()
+        membership2 = CohortMembership(course_user_group=self.cohort_2, user=self.student_2)
+        membership2.save()
 
         result = self._cohort_students_and_upload(
             u'username,email,cohort\n'
diff --git a/openedx/core/djangoapps/course_groups/tests/helpers.py b/openedx/core/djangoapps/course_groups/tests/helpers.py
index 639f6dcd17..260485c406 100644
--- a/openedx/core/djangoapps/course_groups/tests/helpers.py
+++ b/openedx/core/djangoapps/course_groups/tests/helpers.py
@@ -11,7 +11,7 @@ from xmodule.modulestore.django import modulestore
 from xmodule.modulestore import ModuleStoreEnum
 
 from ..cohorts import set_course_cohort_settings
-from ..models import CourseUserGroup, CourseCohort, CourseCohortsSettings
+from ..models import CourseUserGroup, CourseCohort, CourseCohortsSettings, CohortMembership
 
 
 class CohortFactory(DjangoModelFactory):
@@ -41,6 +41,15 @@ class CourseCohortFactory(DjangoModelFactory):
     class Meta(object):
         model = CourseCohort
 
+    @post_generation
+    def memberships(self, create, extracted, **kwargs):  # pylint: disable=unused-argument
+        """
+        Returns the memberships linking users to this cohort.
+        """
+        for user in self.course_user_group.users.all():  # pylint: disable=E1101
+            membership = CohortMembership(user=user, course_user_group=self.course_user_group)
+            membership.save()
+
     course_user_group = factory.SubFactory(CohortFactory)
     assignment_type = 'manual'