From dbbc64ff88addf1054f665f241c12e3e746ac519 Mon Sep 17 00:00:00 2001
From: Eric Fischer <efischer@edx.org>
Date: Fri, 23 Oct 2015 16:19:46 -0400
Subject: [PATCH] CohortMembership Race Condition Test

By using the before_after library, we can force a race condition to reliably
occur in the CohortMembership.save() method. This unit test will do just that,
and ensure that our race-condition-handling code functions as it should.
---
 .../course_groups/tests/test_cohorts.py       | 40 +++++++++++++++++++
 requirements/edx/base.txt                     |  1 +
 2 files changed, 41 insertions(+)

diff --git a/openedx/core/djangoapps/course_groups/tests/test_cohorts.py b/openedx/core/djangoapps/course_groups/tests/test_cohorts.py
index 5ef2b0edbc..d72c53a4c0 100644
--- a/openedx/core/djangoapps/course_groups/tests/test_cohorts.py
+++ b/openedx/core/djangoapps/course_groups/tests/test_cohorts.py
@@ -4,6 +4,7 @@ Tests for cohorts
 # pylint: disable=no-member
 import ddt
 from mock import call, patch
+import before_after
 
 from django.contrib.auth.models import User
 from django.db import IntegrityError
@@ -635,6 +636,45 @@ class TestCohorts(ModuleStoreTestCase):
             lambda: cohorts.add_user_to_cohort(first_cohort, "non_existent_username")
         )
 
+    @patch("openedx.core.djangoapps.course_groups.cohorts.tracker")
+    def add_user_to_cohorts_race_condition(self, mock_tracker):
+        """
+        Makes use of before_after to force a race condition, in order to
+        confirm handling of such conditions is done correctly.
+        """
+        course_user = UserFactory(username="Username", email="a@b.com")
+        course = modulestore().get_course(self.toy_course_key)
+        CourseEnrollment.enroll(course_user, self.toy_course_key)
+        first_cohort = CohortFactory(course_id=course.id, name="FirstCohort")
+        second_cohort = CohortFactory(course_id=course.id, name="SecondCohort")
+
+        # This before_after contextmanager allows for reliable reproduction of a race condition.
+        # It will break before the first save() call creates an entry, and then run add_user_to_cohort again.
+        # Because this second call will write before control is returned, the first call will be writing stale data.
+        # This test confirms that the first add_user_to_cohort call can handle this stale read condition properly.
+        # Proper handling is defined as treating calls as sequential, with write time deciding the order.
+        with before_after.before_after(
+            'django.db.models.Model.save',
+            after_ftn=cohorts.add_user_to_cohort(second_cohort, course_user.username),
+            autospec=True
+        ):
+            # This method will read, then break, then try to write stale data.
+            # It should fail at that, then retry with refreshed data
+            cohorts.add_user_to_cohort(first_cohort, course_user.username)
+
+        mock_tracker.emit.assert_any_call(
+            "edx.cohort.user_add_requested",
+            {
+                "user_id": course_user.id,
+                "cohort_id": first_cohort.id,
+                "cohort_name": first_cohort.name,
+                "previous_cohort_id": second_cohort.id,
+                "previous_cohort_name": second_cohort.name,
+            }
+        )
+        # Note that the following get() will fail with MultipleObjectsReturned if race condition is not handled.
+        self.assertEqual(first_cohort.users.get(), course_user)
+
     def test_get_course_cohort_settings(self):
         """
         Test that cohorts.get_course_cohort_settings is working as expected.
diff --git a/requirements/edx/base.txt b/requirements/edx/base.txt
index 3e086ce62d..bd0db987b3 100644
--- a/requirements/edx/base.txt
+++ b/requirements/edx/base.txt
@@ -120,6 +120,7 @@ django_debug_toolbar==1.3.2
 
 # Used for testing
 astroid==1.3.8
+before_after==0.1.3
 bok-choy==0.4.7
 chrono==1.0.2
 coverage==4.0