Update is_experiment_on and add ExperimentWaffleFlag override test

lms/djangoapps/experiments/flags.py

@@ -170,16 +170,19 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
         return self.is_enabled()
 
     def is_experiment_on(self, course_key=None):
-        return super().is_enabled(course_key)
+        # If no course_key is supplied check the global flag irrespective of courses
+        if course_key is None:
+            return super().is_enabled_without_course_context()
+
+        return super().is_enabled(course_key=course_key)
 
     @contextmanager
     def override(self, active=True, bucket=1):  # pylint: disable=arguments-differ
-        from mock import patch
-        if not active:
-            bucket = 0
-        with patch.object(self, 'get_bucket', return_value=bucket):
-            # TODO We can move this import to the top of the file once this code is
-            # not all contained within the __init__ module.
-            from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
-            with override_waffle_flag(self, active):
+        # Let CourseWaffleFlag override the base waffle flag value
+        with super().override(active=active):
+            # Now override the experiment bucket value
+            from mock import patch
+            if not active:
+                bucket = 0
+            with patch.object(self, 'get_bucket', return_value=bucket):
                 yield

lms/djangoapps/experiments/tests/test_flags.py

@@ -105,7 +105,8 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
     def test_bucket_override(self, active, expected_bucket):
         bucket_flag = CourseWaffleFlag('experiments', 'test.0')
         with bucket_flag.override(active=active):
-            self.assertEqual(self.get_bucket(), expected_bucket)
+            self.assertEqual(self.flag.get_bucket(), expected_bucket)
+            self.assertEqual(self.flag.is_experiment_on(), active)
 
     def test_tracking(self):
         # Run twice, with same request
@@ -143,3 +144,15 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
             self.assertEqual(self.flag.is_enabled_without_course_context(), False)
             self.assertEqual(self.flag.is_enabled(self.key), False)
             self.assertEqual(self.flag.is_enabled(), False)
+
+    @ddt.data(
+        (True, 1, 1),
+        (True, 0, 0),
+        (False, 1, 0),  # bucket is always 0 if the experiment is off
+        (False, 0, 0),
+    )
+    @ddt.unpack
+    def test_override_method(self, active, bucket_override, expected_bucket):
+        with self.flag.override(active=active, bucket=bucket_override):
+            self.assertEqual(self.flag.get_bucket(), expected_bucket)
+            self.assertEqual(self.flag.is_experiment_on(), active)