enrollment_end support for experiments

Now in addition to enrollment_start support in ExperimentWaffleFlag,
you can set an enrollment_end date for your experiment. All enrollments
after this date will get the control experience.

lms/djangoapps/experiments/flags.py

@@ -2,6 +2,7 @@
 Feature flag support for experiments
 """
 
+import datetime
 import logging
 from contextlib import contextmanager
 
@@ -54,6 +55,45 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
         request_cache.set(key, value)
         return value
 
+    def _is_enrollment_inside_date_bounds(self, experiment_values, user, course_key):
+        """ Returns True if the user's enrollment (if any) is valid for the configured experiment date range """
+        from student.models import CourseEnrollment
+
+        enrollment_start = experiment_values.get('enrollment_start')
+        enrollment_end = experiment_values.get('enrollment_end')
+        if not enrollment_start and not enrollment_end:
+            return True  # early exit just to avoid any further lookups
+
+        now = datetime.datetime.now(pytz.utc)
+        enrollment = CourseEnrollment.get_enrollment(user, course_key)
+
+        # If the user isn't enrolled, act like they would enroll right now (this keeps the pre-enroll and post-enroll
+        # experiences the same, if they decide to enroll right now)
+        enrollment_creation_date = enrollment.created if enrollment else now
+
+        # Enrollment must be after any enrollment_start date, if specified
+        if enrollment_start:
+            try:
+                start_date = dateutil.parser.parse(enrollment_start).replace(tzinfo=pytz.UTC)
+            except ValueError:
+                log.exception('Could not parse enrollment start date for experiment %d', self.experiment_id)
+                return False
+            if enrollment_creation_date < start_date:
+                return False
+
+        # Enrollment must be before any enrollment_end date, if specified
+        if enrollment_end:
+            try:
+                end_date = dateutil.parser.parse(enrollment_end).replace(tzinfo=pytz.UTC)
+            except ValueError:
+                log.exception('Could not parse enrollment end date for experiment %d', self.experiment_id)
+                return False
+            if enrollment_creation_date >= end_date:
+                return False
+
+        # All good! Either because the key was not set or because the enrollment was valid
+        return True
+
     def get_bucket(self, course_key=None, track=True):
         """
         Return which bucket number the specified user is in.
@@ -64,7 +104,6 @@ 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 student.models import CourseEnrollment
 
         request = get_current_request()
         if not request:
@@ -89,19 +128,11 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
 
         # Check if the enrollment should even be considered (if it started before the experiment wants, we ignore)
         if course_key and self.experiment_id is not None:
-            start_val = ExperimentKeyValue.objects.filter(experiment_id=self.experiment_id, key='enrollment_start')
-            if start_val:
-                try:
-                    start_date = dateutil.parser.parse(start_val.first().value).replace(tzinfo=pytz.UTC)
-                except ValueError:
-                    log.exception('Could not parse enrollment start date for experiment %d', self.experiment_id)
-                    return self._cache_bucket(experiment_name, 0)
-                enrollment = CourseEnrollment.get_enrollment(request.user, course_key)
-                # Only bail if they have an enrollment and it's old -- if they don't have an enrollment, we want to do
-                # normal bucketing -- consider the case where the experiment has bits that show before you enroll. We
-                # want to keep your bucketing stable before and after you do enroll.
-                if enrollment and enrollment.created < start_date:
-                    return self._cache_bucket(experiment_name, 0)
+            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):
+                return self._cache_bucket(experiment_name, 0)
 
         bucket = stable_bucketing_hash_group(experiment_name, self.num_buckets, request.user.username)
 

lms/djangoapps/experiments/tests/test_flags.py

@@ -62,9 +62,10 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
         self.assertEqual(self.get_bucket(active=False), 0)
 
     @ddt.data(
-        ('2012-01-06', None, 1),  # no enrollment (we allow normal bucketing in this case)
-        ('2012-01-06', '2012-01-05', 0),  # enrolled before experiment
-        ('2012-01-06', '2012-01-07', 1),  # enrolled after experiment
+        ('2012-01-06', None, 1),  # no enrollment, but start is in past (we allow normal bucketing in this case)
+        ('9999-01-06', None, 0),  # no enrollment, but start is in future (we give bucket 0 in that case)
+        ('2012-01-06', '2012-01-05', 0),  # enrolled before experiment start
+        ('2012-01-06', '2012-01-07', 1),  # enrolled after experiment start
         (None, '2012-01-07', 1),  # no experiment date
         ('not-a-date', '2012-01-07', 0),  # bad experiment date
     )
@@ -78,6 +79,24 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
             ExperimentKeyValueFactory(experiment_id=0, key='enrollment_start', value=experiment_start)
         self.assertEqual(self.get_bucket(), expected_bucket)
 
+    @ddt.data(
+        ('2012-01-06', None, 0),  # no enrollment, but end is in past (we give bucket 0 in that case)
+        ('9999-01-06', None, 1),  # no enrollment, but end is in future (we allow normal bucketing in this case)
+        ('2012-01-06', '2012-01-05', 1),  # enrolled before experiment end
+        ('2012-01-06', '2012-01-07', 0),  # enrolled after experiment end
+        (None, '2012-01-07', 1),  # no experiment date
+        ('not-a-date', '2012-01-07', 0),  # bad experiment date
+    )
+    @ddt.unpack
+    def test_enrollment_end(self, experiment_end, enrollment_created, expected_bucket):
+        if enrollment_created:
+            enrollment = CourseEnrollmentFactory(user=self.user, course_id='a/b/c')
+            enrollment.created = parser.parse(enrollment_created).replace(tzinfo=pytz.UTC)
+            enrollment.save()
+        if experiment_end:
+            ExperimentKeyValueFactory(experiment_id=0, key='enrollment_end', value=experiment_end)
+        self.assertEqual(self.get_bucket(), expected_bucket)
+
     @ddt.data(
         (True, 0),
         (False, 1),