Deprecate ExperimentWaffleFlag.override method

This is part of a removal of the many override methods of toggle flag/namespace classes. This allows us to remove imports of test modules from production code.
2020-10-23 11:19:41 +02:00
parent 64a162f69c
commit 16ce16e85e
12 changed files with 117 additions and 99 deletions
--- a/common/djangoapps/student/tests/test_receivers.py
+++ b/common/djangoapps/student/tests/test_receivers.py
@@ -1,10 +1,11 @@
 """ Tests for student signal receivers. """

 from lms.djangoapps.courseware.toggles import (
-    REDIRECT_TO_COURSEWARE_MICROFRONTEND,
    COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES,
-    COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES_FIRST_SECTION_CELEBRATION
+    COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES_FIRST_SECTION_CELEBRATION,
+    REDIRECT_TO_COURSEWARE_MICROFRONTEND
 )
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from student.models import CourseEnrollmentCelebration
 from student.tests.factories import CourseEnrollmentFactory
 from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
@@ -14,7 +15,7 @@ class ReceiversTest(SharedModuleStoreTestCase):
    """
    Tests for dashboard utility functions
    """
-    @REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True)
    @COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES.override(active=True)
    @COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES_FIRST_SECTION_CELEBRATION.override(active=True)
    def test_celebration_created(self):
--- a/lms/djangoapps/course_home_api/dates/v1/tests/test_views.py
+++ b/lms/djangoapps/course_home_api/dates/v1/tests/test_views.py
@@ -3,13 +3,14 @@ Tests for Dates Tab API in the Course Home API
 """

 from datetime import datetime
-import ddt

+import ddt
 from django.urls import reverse

 from course_modes.models import CourseMode
 from lms.djangoapps.course_home_api.tests.utils import BaseCourseHomeTests
 from lms.djangoapps.course_home_api.toggles import COURSE_HOME_MICROFRONTEND, COURSE_HOME_MICROFRONTEND_DATES_TAB
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from openedx.features.content_type_gating.models import ContentTypeGatingConfig
 from student.models import CourseEnrollment

@@ -24,7 +25,7 @@ class DatesTabTestViews(BaseCourseHomeTests):
        self.url = reverse('course-home-dates-tab', args=[self.course.id])
        ContentTypeGatingConfig.objects.create(enabled=True, enabled_as_of=datetime(2017, 1, 1))

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
    @ddt.data(CourseMode.AUDIT, CourseMode.VERIFIED)
    def test_get_authenticated_enrolled_user(self, enrollment_mode):
@@ -37,28 +38,28 @@ class DatesTabTestViews(BaseCourseHomeTests):
        self.assertEqual(response.data.get('learner_is_full_access'), enrollment_mode == CourseMode.VERIFIED)
        self.assertTrue(all(block.get('learner_has_access') for block in date_blocks))

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
    def test_get_authenticated_user_not_enrolled(self):
        response = self.client.get(self.url)
        self.assertEqual(response.status_code, 200)
        self.assertFalse(response.data.get('learner_is_full_access'))

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
    def test_get_unauthenticated_user(self):
        self.client.logout()
        response = self.client.get(self.url)
        self.assertEqual(response.status_code, 401)

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
    def test_get_unknown_course(self):
        url = reverse('course-home-dates-tab', args=['course-v1:unknown+course+2T2020'])
        response = self.client.get(url)
        self.assertEqual(response.status_code, 404)

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
    def test_banner_data_is_returned(self):
        response = self.client.get(self.url)
@@ -68,7 +69,7 @@ class DatesTabTestViews(BaseCourseHomeTests):
        self.assertContains(response, 'content_type_gating_enabled')
        self.assertContains(response, 'verified_upgrade_link')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
    def test_masquerade(self):
        self.switch_to_staff()
--- a/lms/djangoapps/course_home_api/outline/v1/tests/test_views.py
+++ b/lms/djangoapps/course_home_api/outline/v1/tests/test_views.py
@@ -13,6 +13,7 @@ from mock import Mock, patch
 from course_modes.models import CourseMode
 from lms.djangoapps.course_home_api.tests.utils import BaseCourseHomeTests
 from lms.djangoapps.course_home_api.toggles import COURSE_HOME_MICROFRONTEND, COURSE_HOME_MICROFRONTEND_OUTLINE_TAB
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from openedx.core.djangoapps.user_api.preferences.api import set_user_preference
 from openedx.core.djangoapps.user_api.tests.factories import UserCourseTagFactory
 from openedx.features.course_experience import COURSE_ENABLE_UNENROLLED_ACCESS_FLAG, ENABLE_COURSE_GOALS
@@ -33,7 +34,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.url = reverse('course-home-outline-tab', args=[self.course.id])

    @ENABLE_COURSE_GOALS.override(active=True)
-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    @ddt.data(CourseMode.AUDIT, CourseMode.VERIFIED)
    def test_get_authenticated_enrolled_user(self, enrollment_mode):
@@ -66,7 +67,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        if resume_course_url:
            self.assertIn('http://', resume_course_url)

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    def test_get_authenticated_user_not_enrolled(self):
        response = self.client.get(self.url)
@@ -84,14 +85,14 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.assertTrue(all((block.get('title') != "") for block in date_blocks))
        self.assertTrue(all(block.get('date') for block in date_blocks))

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    def test_get_unauthenticated_user(self):
        self.client.logout()
        response = self.client.get(self.url)
        self.assertEqual(response.status_code, 403)

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    def test_masquerade(self):
        user = UserFactory()
@@ -107,7 +108,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.update_masquerade(username=user.username)
        self.assertEqual(self.client.get(self.url).data['dates_widget']['user_timezone'], 'Asia/Tokyo')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    @COURSE_ENABLE_UNENROLLED_ACCESS_FLAG.override()
    def test_handouts(self):
@@ -115,14 +116,14 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.store.create_item(self.user.id, self.course.id, 'course_info', 'handouts', fields={'data': '<p>Hi</p>'})
        self.assertEqual(self.client.get(self.url).data['handouts_html'], '<p>Hi</p>')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    def test_get_unknown_course(self):
        url = reverse('course-home-outline-tab', args=['course-v1:unknown+course+2T2020'])
        response = self.client.get(url)
        self.assertEqual(response.status_code, 404)

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=False)
    @ddt.data(CourseMode.AUDIT, CourseMode.VERIFIED)
    def test_waffle_flag_disabled(self, enrollment_mode):
@@ -130,7 +131,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        response = self.client.get(self.url)
        self.assertEqual(response.status_code, 404)

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    @ddt.data(True, False)
    def test_welcome_message(self, welcome_message_is_dismissed):
@@ -157,14 +158,14 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        welcome_message_html = self.client.get(self.url).data['welcome_message_html']
        self.assertEqual(welcome_message_html, None if welcome_message_is_dismissed else '<p>Welcome</p>')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    @patch('lms.djangoapps.course_home_api.outline.v1.views.generate_offer_html', new=Mock(return_value='<p>Offer</p>'))
    def test_offer_html(self):
        CourseEnrollment.enroll(self.user, self.course.id)
        self.assertEqual(self.client.get(self.url).data['offer_html'], '<p>Offer</p>')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    @patch('lms.djangoapps.course_home_api.outline.v1.views.generate_course_expired_message', new=Mock(return_value='<p>Expired</p>'))
    def test_course_expired_html(self):
@@ -172,7 +173,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.assertEqual(self.client.get(self.url).data['course_expired_html'], '<p>Expired</p>')

    @ENABLE_COURSE_GOALS.override(active=True)
-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    def test_post_course_goal(self):
        CourseEnrollment.enroll(self.user, self.course.id, CourseMode.AUDIT)
@@ -192,7 +193,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.assertIsNotNone(selected_goal)
        self.assertEqual(selected_goal['key'], 'certify')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    @patch.dict('django.conf.settings.FEATURES', {'ENABLE_SPECIAL_EXAMS': True})
    @patch('lms.djangoapps.course_api.blocks.transformers.milestones.get_attempt_status_summary')
@@ -236,7 +237,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.assertIsNotNone(exam_data['due'])
        self.assertEqual(exam_data['icon'], 'fa-foo-bar')

-    @COURSE_HOME_MICROFRONTEND.override(active=True)
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
    def test_assignment(self):
        course = CourseFactory.create()
@@ -267,7 +268,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
        self.assertEqual(ungraded_data['display_name'], 'Ungraded')
        self.assertIsNone(ungraded_data['icon'])

-    @COURSE_HOME_MICROFRONTEND.override()
+    @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override()
    @COURSE_ENABLE_UNENROLLED_ACCESS_FLAG.override()
    @patch('lms.djangoapps.course_home_api.outline.v1.views.generate_offer_html', new=Mock(return_value='<p>Offer</p>'))
--- a/lms/djangoapps/courseware/tests/test_date_summary.py
+++ b/lms/djangoapps/courseware/tests/test_date_summary.py
@@ -10,12 +10,14 @@ import waffle
 from django.contrib.messages.middleware import MessageMiddleware
 from django.test import RequestFactory
 from django.urls import reverse
-from freezegun import freeze_time
 from mock import patch
 from pytz import utc

 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
+from freezegun import freeze_time
+from lms.djangoapps.commerce.models import CommerceConfiguration
+from lms.djangoapps.course_home_api.toggles import COURSE_HOME_MICROFRONTEND, COURSE_HOME_MICROFRONTEND_DATES_TAB
 from lms.djangoapps.courseware.courses import get_course_date_blocks
 from lms.djangoapps.courseware.date_summary import (
    CertificateAvailableDate,
@@ -32,8 +34,7 @@ from lms.djangoapps.courseware.models import (
    DynamicUpgradeDeadlineConfiguration,
    OrgDynamicUpgradeDeadlineConfiguration
 )
-from lms.djangoapps.commerce.models import CommerceConfiguration
-from lms.djangoapps.course_home_api.toggles import COURSE_HOME_MICROFRONTEND, COURSE_HOME_MICROFRONTEND_DATES_TAB
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from lms.djangoapps.verify_student.models import VerificationDeadline
 from lms.djangoapps.verify_student.tests.factories import SoftwareSecurePhotoVerificationFactory
 from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
@@ -44,7 +45,10 @@ from openedx.core.djangoapps.user_api.preferences.api import set_user_preference
 from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.course_duration_limits.models import CourseDurationLimitConfig
 from openedx.features.course_experience import (
-    RELATIVE_DATES_FLAG, DISABLE_UNIFIED_COURSE_TAB_FLAG, UPGRADE_DEADLINE_MESSAGE, CourseHomeMessages
+    DISABLE_UNIFIED_COURSE_TAB_FLAG,
+    RELATIVE_DATES_FLAG,
+    UPGRADE_DEADLINE_MESSAGE,
+    CourseHomeMessages
 )
 from student.tests.factories import TEST_PASSWORD, CourseEnrollmentFactory, UserFactory
 from xmodule.modulestore import ModuleStoreEnum
@@ -145,7 +149,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
        CourseEnrollmentFactory(course_id=course.id, user=user, mode=CourseMode.VERIFIED)
        self.assert_block_types(course, user, expected_blocks)

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_enabled_block_types_with_assignments(self):  # pylint: disable=too-many-statements
        """
        Creates a course with multiple subsections to test all of the different
@@ -300,7 +304,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
                for html_tag in assignment_title_html:
                    self.assertIn(html_tag, assignment_title)

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    @ddt.data(
        ([], 3),
        ([{
@@ -366,7 +370,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
        blocks = get_course_date_blocks(course, user, request, include_past_dates=True)
        self.assertEqual(len(blocks), date_block_count)

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_enabled_block_types_with_expired_course(self):
        course = create_course_run(days_till_start=-100)
        user = create_user()
@@ -543,7 +547,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
        {'weeks_to_complete': 7},  # Weeks to complete > time til end (end date shown)
        {'weeks_to_complete': 4},  # Weeks to complete < time til end (end date not shown)
    )
-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_course_end_date_self_paced(self, cr_details):
        """
        In self-paced courses, the end date will now only show up if the learner
@@ -733,7 +737,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
    )
    @ddt.unpack
    @override_waffle_flag(DISABLE_UNIFIED_COURSE_TAB_FLAG, active=False)
-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_dates_tab_link_render(self, url_name, mfe_active):
        """ The dates tab link should only show for enrolled or staff users """
        course = create_course_run()
@@ -751,7 +755,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
        def assert_html_elements(assert_function, user):
            self.client.login(username=user.username, password=TEST_PASSWORD)
            if mfe_active:
-                with COURSE_HOME_MICROFRONTEND.override(active=True), \
+                with override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True), \
                     COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True):
                    response = self.client.get(url, follow=True)
            else:
--- a/lms/djangoapps/courseware/tests/test_views.py
+++ b/lms/djangoapps/courseware/tests/test_views.py
@@ -11,17 +11,10 @@ import unittest
 from datetime import datetime, timedelta
 from uuid import uuid4

+import ddt
 import six
-from markupsafe import escape
-from mock import MagicMock, PropertyMock, call, create_autospec, patch
-from pytz import UTC, utc
-from six import text_type
-from six.moves import range
-from six.moves.urllib.parse import quote, urlencode
-
 from completion.test_utils import CompletionWaffleTestMixin
 from crum import set_current_request
-import ddt
 from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
 from django.http import Http404, HttpResponseBadRequest
@@ -30,17 +23,24 @@ from django.test.client import Client
 from django.test.utils import override_settings
 from django.urls import reverse, reverse_lazy
 from freezegun import freeze_time
+from markupsafe import escape
 from milestones.tests.utils import MilestonesTestCaseMixin
+from mock import MagicMock, PropertyMock, call, create_autospec, patch
 from opaque_keys.edx.keys import CourseKey, UsageKey
 from opaque_keys.edx.locator import BlockUsageLocator, CourseLocator
+from pytz import UTC, utc
+from six import text_type
+from six.moves import range
+from six.moves.urllib.parse import quote, urlencode
 from web_fragments.fragment import Fragment
 from xblock.core import XBlock
 from xblock.fields import Scope, String
-import lms.djangoapps.courseware.views.views as views

+import lms.djangoapps.courseware.views.views as views
 from capa.tests.response_xml_factory import MultipleChoiceResponseXMLFactory
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.certificates import api as certs_api
 from lms.djangoapps.certificates.models import (
    CertificateGenerationConfiguration,
@@ -63,6 +63,7 @@ from lms.djangoapps.courseware.toggles import (
 from lms.djangoapps.courseware.url_helpers import get_microfrontend_url, get_redirect_url
 from lms.djangoapps.courseware.user_state_client import DjangoXBlockUserStateClient
 from lms.djangoapps.courseware.views.index import show_courseware_mfe_link
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from lms.djangoapps.grades.config.waffle import ASSUME_ZERO_GRADE_IF_ABSENT
 from lms.djangoapps.grades.config.waffle import waffle as grades_waffle
 from lms.djangoapps.verify_student.models import VerificationDeadline
@@ -1445,7 +1446,7 @@ class ProgressPageTests(ProgressPageBaseTests):
    def test_progress_queries(self, enable_waffle, initial, subsequent):
        ContentTypeGatingConfig.objects.create(enabled=True, enabled_as_of=datetime(2018, 1, 1))
        self.setup_course()
-        with grades_waffle().override(ASSUME_ZERO_GRADE_IF_ABSENT, active=enable_waffle):
+        with override_waffle_flag(ASSUME_ZERO_GRADE_IF_ABSENT, active=enable_waffle):
            with self.assertNumQueries(
                initial, table_blacklist=QUERY_COUNT_TABLE_BLACKLIST
            ), check_mongo_calls(1):
@@ -3135,7 +3136,7 @@ class DatesTabTestCase(ModuleStoreTestCase):
        response = self._get_response(self.course)
        self.assertEqual(response.status_code, 200)

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    @patch('edx_django_utils.monitoring.set_custom_attribute')
    def test_defaults(self, mock_set_custom_attribute):
        enrollment = CourseEnrollmentFactory(course_id=self.course.id, user=self.user, mode=CourseMode.VERIFIED)
@@ -3196,7 +3197,7 @@ class DatesTabTestCase(ModuleStoreTestCase):
            # Make sure the assignment type is rendered
            self.assertContains(response, 'Homework:')

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_reset_deadlines_banner_displays(self):
        CourseEnrollmentFactory(course_id=self.course.id, user=self.user, mode=CourseMode.VERIFIED)
        now = datetime.now(utc)
@@ -3250,7 +3251,7 @@ class TestShowCoursewareMFE(TestCase):
        )
        for user, course_key, is_course_staff, preview_active, redirect_active in combos:
            with override_waffle_flag(COURSEWARE_MICROFRONTEND_COURSE_TEAM_PREVIEW, preview_active):
-                with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=redirect_active):
+                with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=redirect_active):
                    assert show_courseware_mfe_link(user, is_course_staff, course_key) is False

    @patch.dict(settings.FEATURES, {'ENABLE_COURSEWARE_MICROFRONTEND': True})
@@ -3270,14 +3271,14 @@ class TestShowCoursewareMFE(TestCase):
        )
        for user, is_course_staff, preview_active, redirect_active in old_mongo_combos:
            with override_waffle_flag(COURSEWARE_MICROFRONTEND_COURSE_TEAM_PREVIEW, preview_active):
-                with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=redirect_active):
+                with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=redirect_active):
                    assert show_courseware_mfe_link(user, is_course_staff, old_course_key) is False

        # We've checked all old-style course keys now, so we can test only the
        # new ones going forward. Now we check combinations of waffle flags and
        # user permissions...
        with override_waffle_flag(COURSEWARE_MICROFRONTEND_COURSE_TEAM_PREVIEW, True):
-            with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True):
+            with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True):
                # (preview=on, redirect=on)
                # Global and Course Staff can see the link.
                self.assertTrue(show_courseware_mfe_link(global_staff_user, True, new_course_key))
@@ -3286,7 +3287,7 @@ class TestShowCoursewareMFE(TestCase):

                # Regular users don't see the link.
                self.assertFalse(show_courseware_mfe_link(regular_user, False, new_course_key))
-            with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=False):
+            with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=False):
                # (preview=on, redirect=off)
                # Global and Course Staff can see the link.
                self.assertTrue(show_courseware_mfe_link(global_staff_user, True, new_course_key))
@@ -3297,7 +3298,7 @@ class TestShowCoursewareMFE(TestCase):
                self.assertFalse(show_courseware_mfe_link(regular_user, False, new_course_key))

        with override_waffle_flag(COURSEWARE_MICROFRONTEND_COURSE_TEAM_PREVIEW, False):
-            with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True):
+            with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True):
                # (preview=off, redirect=on)
                # Global staff see the link anyway
                self.assertTrue(show_courseware_mfe_link(global_staff_user, True, new_course_key))
@@ -3309,7 +3310,7 @@ class TestShowCoursewareMFE(TestCase):

                # Regular users don't see the link.
                self.assertFalse(show_courseware_mfe_link(regular_user, False, new_course_key))
-            with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=False):
+            with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=False):
                # (preview=off, redirect=off)
                # Global staff see the link anyway
                self.assertTrue(show_courseware_mfe_link(global_staff_user, True, new_course_key))
@@ -3368,7 +3369,7 @@ class MFERedirectTests(BaseViewsTestCase):
        # learners will be redirected when the waffle flag is set
        lms_url, mfe_url = self._get_urls()

-        with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True):
+        with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True):
            assert self.client.get(lms_url).url == mfe_url

    def test_staff_no_redirect(self):
@@ -3380,14 +3381,14 @@ class MFERedirectTests(BaseViewsTestCase):
        self.client.login(username=course_staff.username, password='test')

        assert self.client.get(lms_url).status_code == 200
-        with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True):
+        with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True):
            assert self.client.get(lms_url).status_code == 200

        # global staff will never be redirected
        self._create_global_staff_user()
        assert self.client.get(lms_url).status_code == 200

-        with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True):
+        with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True):
            assert self.client.get(lms_url).status_code == 200

    def test_exam_no_redirect(self):
@@ -3397,5 +3398,5 @@ class MFERedirectTests(BaseViewsTestCase):

        lms_url, mfe_url = self._get_urls()

-        with REDIRECT_TO_COURSEWARE_MICROFRONTEND.override(active=True):
+        with override_experiment_waffle_flag(REDIRECT_TO_COURSEWARE_MICROFRONTEND, active=True):
            assert self.client.get(lms_url).status_code == 200
--- a/lms/djangoapps/experiments/flags.py
+++ b/lms/djangoapps/experiments/flags.py
@@ -49,18 +49,18 @@ class ExperimentWaffleFlag(CourseWaffleFlag):

    When writing tests involving an ExperimentWaffleFlag you must not use the
    override_waffle_flag utility. That will only turn the experiment on or off and won't
-    override bucketing. Instead use ExperimentWaffleFlag's override method which
+    override bucketing. Instead use override_experiment_waffle_flag function which
    will do both. Example:

-        with MY_EXPERIMENT_WAFFLE_FLAG.override(active=True, bucket=1):
+        from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
+        with @override_experiment_waffle_flag(MY_EXPERIMENT_WAFFLE_FLAG, active=True, bucket=1):
            ...

    or as a decorator:

-        @MY_EXPERIMENT_WAFFLE_FLAG.override(active=True, bucket=1)
+        @override_experiment_waffle_flag(MY_EXPERIMENT_WAFFLE_FLAG, active=True, bucket=1)
        def test_my_experiment(self):
            ...
-
    """
    def __init__(
            self,
@@ -259,14 +259,3 @@ class ExperimentWaffleFlag(CourseWaffleFlag):
        This disregards `.bucket_flags`.
        """
        return super().is_enabled(course_key)
-
-    @contextmanager
-    def override(self, active=True, bucket=1):  # pylint: disable=arguments-differ
-        # 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
--- a/lms/djangoapps/experiments/tests/test_flags.py
+++ b/lms/djangoapps/experiments/tests/test_flags.py
@@ -2,9 +2,8 @@
 Tests for experimentation feature flags
 """

-import pytz
-
 import ddt
+import pytz
 from crum import set_current_request
 from dateutil import parser
 from django.test.client import RequestFactory
@@ -14,6 +13,7 @@ from opaque_keys.edx.keys import CourseKey

 from experiments.factories import ExperimentKeyValueFactory
 from experiments.flags import ExperimentWaffleFlag
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteFactory
 from openedx.core.djangoapps.waffle_utils import CourseWaffleFlag
 from openedx.core.djangoapps.waffle_utils.models import WaffleFlagCourseOverrideModel
@@ -47,7 +47,7 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
        self.addCleanup(RequestCache.clear_all_namespaces)

    def get_bucket(self, track=False, active=True):
-        # Does not use ExperimentWaffleFlag.override, since that shortcuts get_bucket and we want to test internals
+        # Does not use override_experiment_waffle_flag, since that shortcuts get_bucket and we want to test internals
        with override_waffle_flag(self.flag, active):
            with override_waffle_flag(self.flag.bucket_flags[1], True):
                return self.flag.get_bucket(course_key=self.key, track=track)
@@ -152,7 +152,7 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
    @ddt.unpack
    # Test the override method
    def test_override_method(self, active, bucket_override, expected_bucket):
-        with self.flag.override(active=active, bucket=bucket_override):
+        with override_experiment_waffle_flag(self.flag, active=active, bucket=bucket_override):
            self.assertEqual(self.flag.get_bucket(), expected_bucket)
            self.assertEqual(self.flag.is_experiment_on(), active)

--- a/lms/djangoapps/experiments/testutils.py
+++ b/lms/djangoapps/experiments/testutils.py
@@ -0,0 +1,17 @@
+from contextlib import contextmanager
+
+from mock import patch
+
+from edx_toggles.toggles.testutils import override_waffle_flag
+
+
+@contextmanager
+def override_experiment_waffle_flag(flag, active=True, bucket=1):
+    """
+    Override both the base waffle flag and the experiment bucket value.
+    """
+    if not active:
+        bucket = 0
+    with override_waffle_flag(flag, active):
+        with patch.object(flag, "get_bucket", return_value=bucket):
+            yield
--- a/lms/djangoapps/instructor/tests/test_api.py
+++ b/lms/djangoapps/instructor/tests/test_api.py
@@ -36,6 +36,9 @@ from testfixtures import LogCapture
 from bulk_email.models import BulkEmailFlag, CourseEmail, CourseEmailTemplate
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
+from lms.djangoapps.certificates.api import generate_user_certificates
+from lms.djangoapps.certificates.models import CertificateStatuses
+from lms.djangoapps.certificates.tests.factories import GeneratedCertificateFactory
 from lms.djangoapps.courseware.models import StudentModule
 from lms.djangoapps.courseware.tests.factories import (
    BetaTesterFactory,
@@ -45,9 +48,7 @@ from lms.djangoapps.courseware.tests.factories import (
    UserProfileFactory
 )
 from lms.djangoapps.courseware.tests.helpers import LoginEnrollmentTestCase
-from lms.djangoapps.certificates.api import generate_user_certificates
-from lms.djangoapps.certificates.models import CertificateStatuses
-from lms.djangoapps.certificates.tests.factories import GeneratedCertificateFactory
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from lms.djangoapps.instructor.tests.utils import FakeContentTask, FakeEmail, FakeEmailInfo
 from lms.djangoapps.instructor.views.api import (
    _split_input_list,
@@ -87,8 +88,11 @@ from student.models import (
    get_retired_username_by_username
 )
 from student.roles import (
-    CourseBetaTesterRole, CourseDataResearcherRole, CourseFinanceAdminRole,
-    CourseInstructorRole, CourseSalesAdminRole
+    CourseBetaTesterRole,
+    CourseDataResearcherRole,
+    CourseFinanceAdminRole,
+    CourseInstructorRole,
+    CourseSalesAdminRole
 )
 from student.tests.factories import AdminFactory, UserFactory
 from xmodule.fields import Date
@@ -3997,7 +4001,7 @@ class TestDueDateExtensions(SharedModuleStoreTestCase, LoginEnrollmentTestCase):
            get_extended_due(self.course, self.week3, self.user1)
        )

-    @RELATIVE_DATES_FLAG.override(True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_reset_date(self):
        self.test_change_due_date()
        url = reverse('reset_due_date', kwargs={'course_id': text_type(self.course.id)})
@@ -4011,7 +4015,7 @@ class TestDueDateExtensions(SharedModuleStoreTestCase, LoginEnrollmentTestCase):
            get_extended_due(self.course, self.week1, self.user1)
        )

-    @RELATIVE_DATES_FLAG.override(True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_reset_date_only_in_edx_when(self):
        # Start with a unit that only has a date in edx-when
        self.assertEqual(get_date_for_block(self.course, self.week3, self.user1), None)
@@ -4144,7 +4148,7 @@ class TestDueDateExtensionsDeletedDate(ModuleStoreTestCase, LoginEnrollmentTestC
        self.client.login(username=self.instructor.username, password='test')
        extract_dates(None, self.course.id)

-    @RELATIVE_DATES_FLAG.override(True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_reset_extension_to_deleted_date(self):
        """
        Test that we can delete a due date extension after deleting the normal
--- a/openedx/features/calendar_sync/tests/test_plugins.py
+++ b/openedx/features/calendar_sync/tests/test_plugins.py
@@ -7,12 +7,12 @@ import crum
 import ddt
 from django.test import RequestFactory

-from xmodule.modulestore.tests.django_utils import CourseUserType, SharedModuleStoreTestCase
-from xmodule.modulestore.tests.factories import CourseFactory
-
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.calendar_sync.plugins import CalendarSyncToggleTool
 from openedx.features.course_experience import CALENDAR_SYNC_FLAG, RELATIVE_DATES_FLAG
+from xmodule.modulestore.tests.django_utils import CourseUserType, SharedModuleStoreTestCase
+from xmodule.modulestore.tests.factories import CourseFactory


@ddt.ddt
@@ -35,7 +35,7 @@ class TestCalendarSyncToggleTool(SharedModuleStoreTestCase):
    )
    @ddt.unpack
    @override_waffle_flag(CALENDAR_SYNC_FLAG, active=True)
-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_calendar_sync_toggle_tool_is_enabled(self, user_type, should_be_enabled):
        request = RequestFactory().request()
        request.user = self.create_user_for_course(self.course, user_type)
--- a/openedx/features/course_experience/tests/views/test_course_outline.py
+++ b/openedx/features/course_experience/tests/views/test_course_outline.py
@@ -13,7 +13,7 @@ from completion import waffle
 from completion.models import BlockCompletion
 from completion.test_utils import CompletionWaffleTestMixin
 from django.contrib.sites.models import Site
-from django.test import override_settings, RequestFactory
+from django.test import RequestFactory, override_settings
 from django.urls import reverse
 from django.utils import timezone
 from milestones.tests.utils import MilestonesTestCaseMixin
@@ -27,16 +27,17 @@ from waffle.testutils import override_switch
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
 from gating import api as lms_gating_api
+from lms.djangoapps.course_api.blocks.transformers.milestones import MilestonesAndSpecialExamsTransformer
 from lms.djangoapps.courseware.tests.factories import StaffFactory
 from lms.djangoapps.courseware.tests.helpers import MasqueradeMixin
-from lms.djangoapps.course_api.blocks.transformers.milestones import MilestonesAndSpecialExamsTransformer
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from lms.urls import RESET_COURSE_DEADLINES_NAME
+from openedx.core.djangoapps.course_date_signals.models import SelfPacedRelativeDatesConfig
 from openedx.core.djangoapps.schedules.models import Schedule
 from openedx.core.djangoapps.schedules.tests.factories import ScheduleFactory
-from openedx.core.djangoapps.course_date_signals.models import SelfPacedRelativeDatesConfig
 from openedx.core.lib.gating import api as gating_api
-from openedx.features.course_experience import RELATIVE_DATES_FLAG
 from openedx.features.content_type_gating.models import ContentTypeGatingConfig
+from openedx.features.course_experience import RELATIVE_DATES_FLAG
 from openedx.features.course_experience.views.course_outline import (
    DEFAULT_COMPLETION_TRACKING_START,
    CourseOutlineFragmentView
@@ -48,7 +49,6 @@ from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
 from xmodule.modulestore.tests.factories import CourseFactory, ItemFactory

 from ...utils import get_course_outline_block_tree
-
 from .test_course_home import course_home_url

 TEST_PASSWORD = 'test'
@@ -141,7 +141,7 @@ class TestCourseOutlinePage(SharedModuleStoreTestCase, MasqueradeMixin):
        super(TestCourseOutlinePage, self).setUp()
        self.client.login(username=self.user.username, password=TEST_PASSWORD)

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_outline_details(self):
        for course in self.courses:

@@ -196,7 +196,7 @@ class TestCourseOutlinePage(SharedModuleStoreTestCase, MasqueradeMixin):
        self.assertRegex(content, sequential2.display_name + r'\s*\(1 Question\)\s*</h4>')
        self.assertRegex(content, sequential3.display_name + r'\s*\(2 Questions\)\s*</h4>')

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    @ddt.data(
        ([CourseMode.AUDIT, CourseMode.VERIFIED], CourseMode.AUDIT, False, True),
        ([CourseMode.AUDIT, CourseMode.VERIFIED], CourseMode.VERIFIED, False, True),
@@ -236,7 +236,7 @@ class TestCourseOutlinePage(SharedModuleStoreTestCase, MasqueradeMixin):
        else:
            self.assertNotContains(response, '<div class="banner-cta-text"')

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_reset_course_deadlines(self):
        course = self.courses[0]
        enrollment = CourseEnrollment.objects.get(course_id=course.id)
@@ -263,7 +263,7 @@ class TestCourseOutlinePage(SharedModuleStoreTestCase, MasqueradeMixin):
        updated_staff_schedule = Schedule.objects.get(id=staff_schedule.id)
        self.assertEqual(updated_staff_schedule.start_date, staff_schedule.start_date)

-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_reset_course_deadlines_masquerade_generic_student(self):
        course = self.courses[0]

--- a/openedx/features/personalized_learner_schedules/show_answer/tests/test_show_answer_override.py
+++ b/openedx/features/personalized_learner_schedules/show_answer/tests/test_show_answer_override.py
@@ -1,7 +1,6 @@
 """Tests for Show Answer overrides for self-paced courses."""

 import ddt
-
 from django.test import RequestFactory
 from django.test.utils import override_settings

@@ -9,6 +8,7 @@ from common.lib.xmodule.xmodule.capa_base import SHOWANSWER
 from lms.djangoapps.ccx.tests.test_overrides import inject_field_overrides
 from lms.djangoapps.courseware.model_data import FieldDataCache
 from lms.djangoapps.courseware.module_render import get_module
+from lms.djangoapps.experiments.testutils import override_experiment_waffle_flag
 from openedx.features.course_experience import RELATIVE_DATES_FLAG
 from xmodule.modulestore.tests.django_utils import ModuleStoreTestCase
 from xmodule.modulestore.tests.factories import CourseFactory
@@ -36,7 +36,7 @@ class ShowAnswerFieldOverrideTest(ModuleStoreTestCase):

    @ddt.data(True, False)
    def test_override_enabled_for(self, active):
-        with RELATIVE_DATES_FLAG.override(active=active):
+        with override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=active):
            # Instructor paced course will just have the default value
            ip_course = self.setup_course()
            course_module = self.get_course_module(ip_course)
@@ -63,7 +63,7 @@ class ShowAnswerFieldOverrideTest(ModuleStoreTestCase):
        (SHOWANSWER.ALWAYS, SHOWANSWER.ALWAYS),
    )
    @ddt.unpack
-    @RELATIVE_DATES_FLAG.override(active=True)
+    @override_experiment_waffle_flag(RELATIVE_DATES_FLAG, active=True)
    def test_get(self, initial_value, expected_final_value):
        course = self.setup_course(self_paced=True, showanswer=initial_value)
        course_module = self.get_course_module(course)