diff --git a/cms/djangoapps/contentstore/tests/test_course_settings.py b/cms/djangoapps/contentstore/tests/test_course_settings.py
index 814d4f69b1..7cc2455fa4 100644
--- a/cms/djangoapps/contentstore/tests/test_course_settings.py
+++ b/cms/djangoapps/contentstore/tests/test_course_settings.py
@@ -30,8 +30,8 @@ from cms.djangoapps.models.settings.course_metadata import CourseMetadata
 from cms.djangoapps.models.settings.encoder import CourseSettingsEncoder
 from cms.djangoapps.models.settings.waffle import MATERIAL_RECOMPUTE_ONLY_FLAG
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.models.course_details import CourseDetails
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.roles import CourseInstructorRole, CourseStaffRole
 from student.tests.factories import UserFactory
 from util import milestones_helpers
diff --git a/cms/djangoapps/contentstore/views/tests/test_videos.py b/cms/djangoapps/contentstore/views/tests/test_videos.py
index bed9bf82c7..6cf01927b4 100644
--- a/cms/djangoapps/contentstore/views/tests/test_videos.py
+++ b/cms/djangoapps/contentstore/views/tests/test_videos.py
@@ -33,6 +33,7 @@ from waffle.testutils import override_flag
 from cms.djangoapps.contentstore.models import VideoUploadConfig
 from cms.djangoapps.contentstore.tests.utils import CourseTestCase
 from cms.djangoapps.contentstore.utils import reverse_course_url
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.profile_images.tests.helpers import make_image_file
 from openedx.core.djangoapps.video_pipeline.config.waffle import (
     DEPRECATE_YOUTUBE,
@@ -41,7 +42,6 @@ from openedx.core.djangoapps.video_pipeline.config.waffle import (
 )
 from openedx.core.djangoapps.video_pipeline.models import VEMPipelineIntegration
 from openedx.core.djangoapps.waffle_utils.models import WaffleFlagCourseOverrideModel
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from xmodule.modulestore.tests.factories import CourseFactory
 
 from ..videos import (
diff --git a/common/djangoapps/student/tests/test_receivers.py b/common/djangoapps/student/tests/test_receivers.py
index 29a21065f3..40376dc4a6 100644
--- a/common/djangoapps/student/tests/test_receivers.py
+++ b/common/djangoapps/student/tests/test_receivers.py
@@ -1,5 +1,6 @@
 """ Tests for student signal receivers. """
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.courseware.toggles import (
     COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES,
     COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES_FIRST_SECTION_CELEBRATION,
@@ -16,8 +17,8 @@ class ReceiversTest(SharedModuleStoreTestCase):
     Tests for dashboard utility functions
     """
     @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)
+    @override_waffle_flag(COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES, active=True)
+    @override_waffle_flag(COURSEWARE_MICROFRONTEND_PROGRESS_MILESTONES_FIRST_SECTION_CELEBRATION, active=True)
     def test_celebration_created(self):
         """ Test that we make celebration objects when enrollments are created """
         self.assertEqual(CourseEnrollmentCelebration.objects.count(), 0)
diff --git a/common/djangoapps/student/tests/test_views.py b/common/djangoapps/student/tests/test_views.py
index 3096cf336d..d2e229d63f 100644
--- a/common/djangoapps/student/tests/test_views.py
+++ b/common/djangoapps/student/tests/test_views.py
@@ -25,6 +25,7 @@ from pyquery import PyQuery as pq
 from six.moves import range
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from entitlements.tests.factories import CourseEntitlementFactory
 from lms.djangoapps.certificates.tests.factories import GeneratedCertificateFactory
 from openedx.core.djangoapps.catalog.tests.factories import ProgramFactory
@@ -33,7 +34,6 @@ from openedx.core.djangoapps.content.course_overviews.tests.factories import Cou
 from openedx.core.djangoapps.schedules.config import COURSE_UPDATE_WAFFLE_FLAG
 from openedx.core.djangoapps.schedules.tests.factories import ScheduleFactory
 from openedx.core.djangoapps.site_configuration.tests.test_util import with_site_configuration_context
-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.tests.views.helpers import add_course_mode
 from student.helpers import DISABLE_UNENROLL_CERT_STATES
diff --git a/common/lib/xmodule/xmodule/tests/test_sequence.py b/common/lib/xmodule/xmodule/tests/test_sequence.py
index 318d788157..46278f72a5 100644
--- a/common/lib/xmodule/xmodule/tests/test_sequence.py
+++ b/common/lib/xmodule/xmodule/tests/test_sequence.py
@@ -16,7 +16,7 @@ from mock import Mock, patch
 from six.moves import range
 from web_fragments.fragment import Fragment
 
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
+from edx_toggles.toggles.testutils import override_waffle_flag
 from student.tests.factories import UserFactory
 from xmodule.seq_module import TIMED_EXAM_GATING_WAFFLE_FLAG, SequenceModule
 from xmodule.tests import get_test_system
diff --git a/lms/djangoapps/certificates/tests/test_webview_views.py b/lms/djangoapps/certificates/tests/test_webview_views.py
index 1ce54771a9..a3997f33c6 100644
--- a/lms/djangoapps/certificates/tests/test_webview_views.py
+++ b/lms/djangoapps/certificates/tests/test_webview_views.py
@@ -14,7 +14,6 @@ from django.test.client import Client, RequestFactory
 from django.test.utils import override_settings
 from django.urls import reverse
 from mock import patch
-from urllib.parse import urlencode
 
 from course_modes.models import CourseMode
 from lms.djangoapps.badges.events.course_complete import get_completion_badge
diff --git a/lms/djangoapps/course_api/blocks/tests/test_api.py b/lms/djangoapps/course_api/blocks/tests/test_api.py
index a632aef1ff..4b0fe82a42 100644
--- a/lms/djangoapps/course_api/blocks/tests/test_api.py
+++ b/lms/djangoapps/course_api/blocks/tests/test_api.py
@@ -12,7 +12,6 @@ from mock import patch
 
 from openedx.core.djangoapps.content.block_structure.api import clear_course_from_cache
 from openedx.core.djangoapps.content.block_structure.config import STORAGE_BACKING_FOR_CACHE, waffle
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.tests.factories import UserFactory
 from xmodule.modulestore import ModuleStoreEnum
 from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
diff --git a/lms/djangoapps/course_home_api/dates/v1/tests/test_views.py b/lms/djangoapps/course_home_api/dates/v1/tests/test_views.py
index 42d8a6b479..0ad6288e27 100644
--- a/lms/djangoapps/course_home_api/dates/v1/tests/test_views.py
+++ b/lms/djangoapps/course_home_api/dates/v1/tests/test_views.py
@@ -8,6 +8,7 @@ import ddt
 from django.urls import reverse
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 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
@@ -26,7 +27,7 @@ class DatesTabTestViews(BaseCourseHomeTests):
         ContentTypeGatingConfig.objects.create(enabled=True, enabled_as_of=datetime(2017, 1, 1))
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, active=True)
     @ddt.data(CourseMode.AUDIT, CourseMode.VERIFIED)
     def test_get_authenticated_enrolled_user(self, enrollment_mode):
         CourseEnrollment.enroll(self.user, self.course.id, enrollment_mode)
@@ -39,28 +40,28 @@ class DatesTabTestViews(BaseCourseHomeTests):
         self.assertTrue(all(block.get('learner_has_access') for block in date_blocks))
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, 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'))
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, active=True)
     def test_get_unauthenticated_user(self):
         self.client.logout()
         response = self.client.get(self.url)
         self.assertEqual(response.status_code, 401)
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, 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)
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, active=True)
     def test_banner_data_is_returned(self):
         response = self.client.get(self.url)
         self.assertEqual(response.status_code, 200)
@@ -70,7 +71,7 @@ class DatesTabTestViews(BaseCourseHomeTests):
         self.assertContains(response, 'verified_upgrade_link')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, active=True)
     def test_masquerade(self):
         self.switch_to_staff()
         CourseEnrollment.enroll(self.user, self.course.id, 'audit')
diff --git a/lms/djangoapps/course_home_api/outline/v1/tests/test_views.py b/lms/djangoapps/course_home_api/outline/v1/tests/test_views.py
index b5e49cb78d..683c49958d 100644
--- a/lms/djangoapps/course_home_api/outline/v1/tests/test_views.py
+++ b/lms/djangoapps/course_home_api/outline/v1/tests/test_views.py
@@ -11,6 +11,7 @@ from django.urls import reverse
 from mock import Mock, patch
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 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
@@ -33,9 +34,9 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         super().setUp()
         self.url = reverse('course-home-outline-tab', args=[self.course.id])
 
-    @ENABLE_COURSE_GOALS.override(active=True)
+    @override_waffle_flag(ENABLE_COURSE_GOALS, active=True)
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     @ddt.data(CourseMode.AUDIT, CourseMode.VERIFIED)
     def test_get_authenticated_enrolled_user(self, enrollment_mode):
         CourseEnrollment.enroll(self.user, self.course.id, enrollment_mode)
@@ -68,7 +69,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
             self.assertIn('http://', resume_course_url)
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     def test_get_authenticated_user_not_enrolled(self):
         response = self.client.get(self.url)
         self.assertEqual(response.status_code, 200)
@@ -86,14 +87,14 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertTrue(all(block.get('date') for block in date_blocks))
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     def test_get_unauthenticated_user(self):
         self.client.logout()
         response = self.client.get(self.url)
         self.assertEqual(response.status_code, 403)
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     def test_masquerade(self):
         user = UserFactory()
         set_user_preference(user, 'time_zone', 'Asia/Tokyo')
@@ -109,22 +110,22 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertEqual(self.client.get(self.url).data['dates_widget']['user_timezone'], 'Asia/Tokyo')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
-    @COURSE_ENABLE_UNENROLLED_ACCESS_FLAG.override()
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
+    @override_waffle_flag(COURSE_ENABLE_UNENROLLED_ACCESS_FLAG, active=True)
     def test_handouts(self):
         CourseEnrollment.enroll(self.user, self.course.id)
         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>')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, 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)
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=False)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=False)
     @ddt.data(CourseMode.AUDIT, CourseMode.VERIFIED)
     def test_waffle_flag_disabled(self, enrollment_mode):
         CourseEnrollment.enroll(self.user, self.course.id, enrollment_mode)
@@ -132,7 +133,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertEqual(response.status_code, 404)
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     @ddt.data(True, False)
     def test_welcome_message(self, welcome_message_is_dismissed):
         CourseEnrollment.enroll(self.user, self.course.id)
@@ -159,22 +160,22 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertEqual(welcome_message_html, None if welcome_message_is_dismissed else '<p>Welcome</p>')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, 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>')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, 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):
         CourseEnrollment.enroll(self.user, self.course.id)
         self.assertEqual(self.client.get(self.url).data['course_expired_html'], '<p>Expired</p>')
 
-    @ENABLE_COURSE_GOALS.override(active=True)
+    @override_waffle_flag(ENABLE_COURSE_GOALS, active=True)
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     def test_post_course_goal(self):
         CourseEnrollment.enroll(self.user, self.course.id, CourseMode.AUDIT)
 
@@ -194,7 +195,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertEqual(selected_goal['key'], 'certify')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     @patch.dict('django.conf.settings.FEATURES', {'ENABLE_SPECIAL_EXAMS': True})
     @patch('lms.djangoapps.course_api.blocks.transformers.milestones.get_attempt_status_summary')
     def test_proctored_exam(self, mock_summary):
@@ -238,7 +239,7 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertEqual(exam_data['icon'], 'fa-foo-bar')
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override(active=True)
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
     def test_assignment(self):
         course = CourseFactory.create()
         with self.store.bulk_operations(course.id):
@@ -269,8 +270,8 @@ class OutlineTabTestViews(BaseCourseHomeTests):
         self.assertIsNone(ungraded_data['icon'])
 
     @override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True)
-    @COURSE_HOME_MICROFRONTEND_OUTLINE_TAB.override()
-    @COURSE_ENABLE_UNENROLLED_ACCESS_FLAG.override()
+    @override_waffle_flag(COURSE_HOME_MICROFRONTEND_OUTLINE_TAB, active=True)
+    @override_waffle_flag(COURSE_ENABLE_UNENROLLED_ACCESS_FLAG, active=True)
     @patch('lms.djangoapps.course_home_api.outline.v1.views.generate_offer_html', new=Mock(return_value='<p>Offer</p>'))
     @patch('lms.djangoapps.course_home_api.outline.v1.views.generate_course_expired_message', new=Mock(return_value='<p>Expired</p>'))
     @ddt.data(*itertools.product([True, False], [True, False],
diff --git a/lms/djangoapps/course_home_api/progress/v1/tests/test_views.py b/lms/djangoapps/course_home_api/progress/v1/tests/test_views.py
index 93beb59191..9b1f27683a 100644
--- a/lms/djangoapps/course_home_api/progress/v1/tests/test_views.py
+++ b/lms/djangoapps/course_home_api/progress/v1/tests/test_views.py
@@ -3,15 +3,14 @@ Tests for Progress Tab API in the Course Home API
 """
 
 import ddt
-
 from django.urls import reverse
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.course_home_api.tests.utils import BaseCourseHomeTests
 from lms.djangoapps.course_home_api.toggles import COURSE_HOME_MICROFRONTEND
 from lms.djangoapps.verify_student.models import ManualVerification
 from openedx.core.djangoapps.user_api.preferences.api import set_user_preference
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.models import CourseEnrollment
 from student.tests.factories import UserFactory
 
diff --git a/lms/djangoapps/courseware/tests/test_about.py b/lms/djangoapps/courseware/tests/test_about.py
index 660e1dd81c..66177e0d7c 100644
--- a/lms/djangoapps/courseware/tests/test_about.py
+++ b/lms/djangoapps/courseware/tests/test_about.py
@@ -19,8 +19,8 @@ from six import text_type
 from waffle.testutils import override_switch
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.ccx.tests.factories import CcxFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.course_experience import COURSE_ENABLE_UNENROLLED_ACCESS_FLAG
 from openedx.features.course_experience.waffle import ENABLE_COURSE_ABOUT_SIDEBAR_HTML
 from openedx.features.course_experience.waffle import WAFFLE_NAMESPACE as COURSE_EXPERIENCE_WAFFLE_NAMESPACE
diff --git a/lms/djangoapps/courseware/tests/test_course_tools.py b/lms/djangoapps/courseware/tests/test_course_tools.py
index 75b7300786..5c8a0d60c5 100644
--- a/lms/djangoapps/courseware/tests/test_course_tools.py
+++ b/lms/djangoapps/courseware/tests/test_course_tools.py
@@ -12,13 +12,12 @@ from mock import patch
 
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
-from lms.djangoapps.courseware.course_tools import FinancialAssistanceTool
-from lms.djangoapps.courseware.course_tools import VerifiedUpgradeTool
+from edx_toggles.toggles.testutils import override_waffle_flag
+from lms.djangoapps.courseware.course_tools import FinancialAssistanceTool, VerifiedUpgradeTool
 from lms.djangoapps.courseware.models import DynamicUpgradeDeadlineConfiguration
 from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
 from openedx.core.djangoapps.schedules.config import CREATE_SCHEDULE_WAFFLE_FLAG
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
 from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
 from xmodule.modulestore.tests.factories import CourseFactory
diff --git a/lms/djangoapps/courseware/tests/test_date_summary.py b/lms/djangoapps/courseware/tests/test_date_summary.py
index e0d70cabec..9b745050ff 100644
--- a/lms/djangoapps/courseware/tests/test_date_summary.py
+++ b/lms/djangoapps/courseware/tests/test_date_summary.py
@@ -10,6 +10,7 @@ import waffle
 from django.contrib.messages.middleware import MessageMiddleware
 from django.test import RequestFactory
 from django.urls import reverse
+from edx_toggles.toggles.testutils import override_waffle_flag
 from mock import patch
 from pytz import utc
 
@@ -42,7 +43,6 @@ from openedx.core.djangoapps.schedules.signals import CREATE_SCHEDULE_WAFFLE_FLA
 from openedx.core.djangoapps.self_paced.models import SelfPacedConfiguration
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteFactory
 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 (
     DISABLE_UNIFIED_COURSE_TAB_FLAG,
@@ -756,7 +756,7 @@ class CourseDateSummaryTest(SharedModuleStoreTestCase):
             self.client.login(username=user.username, password=TEST_PASSWORD)
             if mfe_active:
                 with override_experiment_waffle_flag(COURSE_HOME_MICROFRONTEND, active=True), \
-                     COURSE_HOME_MICROFRONTEND_DATES_TAB.override(active=True):
+                     override_waffle_flag(COURSE_HOME_MICROFRONTEND_DATES_TAB, active=True):
                     response = self.client.get(url, follow=True)
             else:
                 response = self.client.get(url, follow=True)
diff --git a/lms/djangoapps/courseware/tests/test_entrance_exam.py b/lms/djangoapps/courseware/tests/test_entrance_exam.py
index aa091d8d86..89b7721d9a 100644
--- a/lms/djangoapps/courseware/tests/test_entrance_exam.py
+++ b/lms/djangoapps/courseware/tests/test_entrance_exam.py
@@ -10,6 +10,7 @@ from milestones.tests.utils import MilestonesTestCaseMixin
 from mock import Mock, patch
 
 from capa.tests.response_xml_factory import MultipleChoiceResponseXMLFactory
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.courseware.entrance_exams import (
     course_has_entrance_exam,
     get_entrance_exam_content,
@@ -20,7 +21,6 @@ from lms.djangoapps.courseware.model_data import FieldDataCache
 from lms.djangoapps.courseware.module_render import get_module, handle_xblock_callback, toc_for_course
 from lms.djangoapps.courseware.tests.factories import InstructorFactory, RequestFactoryNoCsrf, StaffFactory, UserFactory
 from lms.djangoapps.courseware.tests.helpers import LoginEnrollmentTestCase
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import get_mock_request
 from openedx.features.course_experience import DISABLE_COURSE_OUTLINE_PAGE_FLAG, DISABLE_UNIFIED_COURSE_TAB_FLAG
 from student.models import CourseEnrollment
diff --git a/lms/djangoapps/courseware/tests/test_masquerade.py b/lms/djangoapps/courseware/tests/test_masquerade.py
index cbb6d81eaf..1fbd29ef8f 100644
--- a/lms/djangoapps/courseware/tests/test_masquerade.py
+++ b/lms/djangoapps/courseware/tests/test_masquerade.py
@@ -19,7 +19,9 @@ from xblock.runtime import DictKeyValueStore
 
 from capa.tests.response_xml_factory import OptionResponseXMLFactory
 from lms.djangoapps.courseware.masquerade import (
-    CourseMasquerade, MasqueradingKeyValueStore, get_masquerading_user_group,
+    CourseMasquerade,
+    MasqueradingKeyValueStore,
+    get_masquerading_user_group
 )
 from lms.djangoapps.courseware.tests.factories import StaffFactory
 from lms.djangoapps.courseware.tests.helpers import LoginEnrollmentTestCase, MasqueradeMixin, masquerade_as_group_member
diff --git a/lms/djangoapps/courseware/tests/test_navigation.py b/lms/djangoapps/courseware/tests/test_navigation.py
index d9b1b55bcc..032eb9733b 100644
--- a/lms/djangoapps/courseware/tests/test_navigation.py
+++ b/lms/djangoapps/courseware/tests/test_navigation.py
@@ -12,9 +12,9 @@ from mock import patch
 from six import text_type
 from six.moves import range
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.courseware.tests.factories import GlobalStaffFactory
 from lms.djangoapps.courseware.tests.helpers import LoginEnrollmentTestCase
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.course_experience import DISABLE_COURSE_OUTLINE_PAGE_FLAG
 from student.tests.factories import UserFactory
 from xmodule.modulestore.django import modulestore
diff --git a/lms/djangoapps/courseware/tests/test_tabs.py b/lms/djangoapps/courseware/tests/test_tabs.py
index 94f67f4669..526bf14483 100644
--- a/lms/djangoapps/courseware/tests/test_tabs.py
+++ b/lms/djangoapps/courseware/tests/test_tabs.py
@@ -13,6 +13,7 @@ from mock import MagicMock, Mock, patch
 from six import text_type
 from six.moves import range
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.courseware.courses import get_course_by_id
 from lms.djangoapps.courseware.tabs import (
     CourseInfoTab,
@@ -26,7 +27,6 @@ from lms.djangoapps.courseware.tabs import (
 from lms.djangoapps.courseware.tests.factories import InstructorFactory, StaffFactory
 from lms.djangoapps.courseware.tests.helpers import LoginEnrollmentTestCase
 from lms.djangoapps.courseware.views.views import StaticCourseTabView, get_static_tab_fragment
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import get_mock_request
 from openedx.features.course_experience import DISABLE_UNIFIED_COURSE_TAB_FLAG
 from student.models import CourseEnrollment
diff --git a/lms/djangoapps/courseware/tests/test_views.py b/lms/djangoapps/courseware/tests/test_views.py
index 350c3bb71c..d058fcde88 100644
--- a/lms/djangoapps/courseware/tests/test_views.py
+++ b/lms/djangoapps/courseware/tests/test_views.py
@@ -74,7 +74,7 @@ from openedx.core.djangoapps.content.course_overviews.models import CourseOvervi
 from openedx.core.djangoapps.crawlers.models import CrawlersConfig
 from openedx.core.djangoapps.credit.api import set_credit_requirements
 from openedx.core.djangoapps.credit.models import CreditCourse, CreditProvider
-from openedx.core.djangoapps.waffle_utils.testutils import WAFFLE_TABLES, override_waffle_flag
+from openedx.core.djangoapps.waffle_utils.testutils import WAFFLE_TABLES
 from openedx.core.djangolib.testing.utils import get_mock_request
 from openedx.core.lib.gating import api as gating_api
 from openedx.core.lib.url_utils import quote_slashes
@@ -1446,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 override_waffle_flag(ASSUME_ZERO_GRADE_IF_ABSENT, active=enable_waffle):
+        with grades_waffle().override(ASSUME_ZERO_GRADE_IF_ABSENT, active=enable_waffle):
             with self.assertNumQueries(
                 initial, table_blacklist=QUERY_COUNT_TABLE_BLACKLIST
             ), check_mongo_calls(1):
diff --git a/lms/djangoapps/experiments/tests/test_flags.py b/lms/djangoapps/experiments/tests/test_flags.py
index 180f309d57..e729c2da35 100644
--- a/lms/djangoapps/experiments/tests/test_flags.py
+++ b/lms/djangoapps/experiments/tests/test_flags.py
@@ -11,13 +11,13 @@ from edx_django_utils.cache import RequestCache
 from mock import patch
 from opaque_keys.edx.keys import CourseKey
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 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
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
 from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
 
@@ -105,7 +105,7 @@ class ExperimentWaffleFlagTests(SharedModuleStoreTestCase):
     @ddt.unpack
     def test_forcing_bucket(self, active, expected_bucket):
         bucket_flag = CourseWaffleFlag('experiments', 'test.0', __name__)
-        with bucket_flag.override(active=active):
+        with override_waffle_flag(bucket_flag, active=active):
             self.assertEqual(self.get_bucket(), expected_bucket)
 
     def test_tracking(self):
diff --git a/lms/djangoapps/experiments/tests/test_views_custom.py b/lms/djangoapps/experiments/tests/test_views_custom.py
index 7a8134cfba..597f8b9991 100644
--- a/lms/djangoapps/experiments/tests/test_views_custom.py
+++ b/lms/djangoapps/experiments/tests/test_views_custom.py
@@ -5,21 +5,19 @@ Tests for experimentation views
 
 from datetime import timedelta
 from uuid import uuid4
-import six
 
+import six
 from django.urls import reverse
 from django.utils.timezone import now
 from rest_framework.test import APITestCase
 
 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.course_blocks.transformers.tests.helpers import ModuleStoreTestCase
-from student.tests.factories import CourseEnrollmentFactory, UserFactory
-
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
-from xmodule.modulestore.tests.factories import CourseFactory
-
 from lms.djangoapps.experiments.views_custom import MOBILE_UPSELL_FLAG
+from student.tests.factories import CourseEnrollmentFactory, UserFactory
+from xmodule.modulestore.tests.factories import CourseFactory
 
 CROSS_DOMAIN_REFERER = 'https://ecommerce.edx.org'
 
diff --git a/lms/djangoapps/grades/config/waffle.py b/lms/djangoapps/grades/config/waffle.py
index 41a171428c..fde2985432 100644
--- a/lms/djangoapps/grades/config/waffle.py
+++ b/lms/djangoapps/grades/config/waffle.py
@@ -84,6 +84,9 @@ def waffle():
 def waffle_flags():
     """
     Returns the namespaced, cached, audited Waffle flags dictionary for Grades.
+
+    WARNING: do not replicate this pattern. Instead of declaring waffle flag names as strings, you should create
+    WaffleFlag and CourseWaffleFlag objects as top-level constants.
     """
     namespace = WaffleFlagNamespace(name=WAFFLE_NAMESPACE, log_prefix=u'Grades: ')
     return {
diff --git a/lms/djangoapps/grades/rest_api/v1/tests/test_gradebook_views.py b/lms/djangoapps/grades/rest_api/v1/tests/test_gradebook_views.py
index 05350fc198..14b46691df 100644
--- a/lms/djangoapps/grades/rest_api/v1/tests/test_gradebook_views.py
+++ b/lms/djangoapps/grades/rest_api/v1/tests/test_gradebook_views.py
@@ -18,6 +18,7 @@ from rest_framework.test import APITestCase
 from six import text_type
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.certificates.models import CertificateStatuses, GeneratedCertificate
 from lms.djangoapps.courseware.tests.factories import InstructorFactory, StaffFactory
 from lms.djangoapps.grades.config.waffle import WRITABLE_GRADEBOOK, waffle_flags
@@ -36,7 +37,6 @@ from lms.djangoapps.grades.rest_api.v1.views import CourseEnrollmentPagination
 from lms.djangoapps.grades.subsection_grade import ReadSubsectionGrade
 from openedx.core.djangoapps.content.course_overviews.tests.factories import CourseOverviewFactory
 from openedx.core.djangoapps.course_groups.tests.helpers import CohortFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
 from xmodule.modulestore.tests.django_utils import TEST_DATA_SPLIT_MODULESTORE, SharedModuleStoreTestCase
 from xmodule.modulestore.tests.factories import CourseFactory, ItemFactory
diff --git a/lms/djangoapps/grades/tests/test_tasks.py b/lms/djangoapps/grades/tests/test_tasks.py
index c836f1e33e..50b7dd087d 100644
--- a/lms/djangoapps/grades/tests/test_tasks.py
+++ b/lms/djangoapps/grades/tests/test_tasks.py
@@ -17,6 +17,7 @@ from django.utils import timezone
 from mock import MagicMock, patch
 from six.moves import range
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.grades import tasks
 from lms.djangoapps.grades.config.models import PersistentGradesEnabledFlag
 from lms.djangoapps.grades.config.waffle import ENFORCE_FREEZE_GRADE_AFTER_COURSE_END, waffle_flags
@@ -33,7 +34,6 @@ from lms.djangoapps.grades.tasks import (
     recalculate_subsection_grade_v3
 )
 from openedx.core.djangoapps.content.block_structure.exceptions import BlockStructureNotFound
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.models import CourseEnrollment, anonymous_id_for_user
 from student.tests.factories import UserFactory
 from track.event_transaction_utils import create_new_event_transaction_id, get_event_transaction_id
diff --git a/lms/djangoapps/instructor/tests/test_api.py b/lms/djangoapps/instructor/tests/test_api.py
index 5420d8edd7..125b45d0e5 100644
--- a/lms/djangoapps/instructor/tests/test_api.py
+++ b/lms/djangoapps/instructor/tests/test_api.py
@@ -36,6 +36,7 @@ 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 edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.certificates.api import generate_user_certificates
 from lms.djangoapps.certificates.models import CertificateStatuses
 from lms.djangoapps.certificates.tests.factories import GeneratedCertificateFactory
@@ -68,7 +69,6 @@ from openedx.core.djangoapps.django_comment_common.utils import seed_permissions
 from openedx.core.djangoapps.schedules.tests.factories import ScheduleFactory
 from openedx.core.djangoapps.site_configuration import helpers as configuration_helpers
 from openedx.core.djangoapps.site_configuration.tests.mixins import SiteMixin
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.lib.teams_config import TeamsConfig
 from openedx.core.lib.xblock_utils import grade_histogram
 from openedx.features.course_experience import RELATIVE_DATES_FLAG
diff --git a/lms/djangoapps/instructor/tests/views/test_instructor_dashboard.py b/lms/djangoapps/instructor/tests/views/test_instructor_dashboard.py
index b221f8b734..19fe30496f 100644
--- a/lms/djangoapps/instructor/tests/views/test_instructor_dashboard.py
+++ b/lms/djangoapps/instructor/tests/views/test_instructor_dashboard.py
@@ -20,15 +20,15 @@ from six.moves import range
 
 from common.test.utils import XssTestMixin
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from edxmako.shortcuts import render_to_response
 from lms.djangoapps.courseware.tabs import get_course_tab_list
 from lms.djangoapps.courseware.tests.factories import StaffFactory, StudentModuleFactory, UserFactory
 from lms.djangoapps.courseware.tests.helpers import LoginEnrollmentTestCase
 from lms.djangoapps.grades.config.waffle import WRITABLE_GRADEBOOK, waffle_flags
-from lms.djangoapps.instructor.views.gradebook_api import calculate_page_info
 from lms.djangoapps.instructor.toggles import DATA_DOWNLOAD_V2
+from lms.djangoapps.instructor.views.gradebook_api import calculate_page_info
 from openedx.core.djangoapps.site_configuration.models import SiteConfiguration
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from student.models import CourseEnrollment
 from student.roles import CourseFinanceAdminRole
 from student.tests.factories import AdminFactory, CourseAccessRoleFactory, CourseEnrollmentFactory
diff --git a/openedx/core/djangoapps/content/learning_sequences/api/tests/test_outlines.py b/openedx/core/djangoapps/content/learning_sequences/api/tests/test_outlines.py
index 814d15fd2f..90e961ce64 100644
--- a/openedx/core/djangoapps/content/learning_sequences/api/tests/test_outlines.py
+++ b/openedx/core/djangoapps/content/learning_sequences/api/tests/test_outlines.py
@@ -10,17 +10,15 @@ from edx_when.api import set_dates_for_course
 from opaque_keys.edx.keys import CourseKey, UsageKey
 from opaque_keys.edx.locator import BlockUsageLocator
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.courseware.tests.factories import BetaTesterFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import CacheIsolationTestCase
 from openedx.features.course_experience import COURSE_ENABLE_UNENROLLED_ACCESS_FLAG
 from student.auth import user_has_role
 from student.models import CourseEnrollment
 from student.roles import CourseBetaTesterRole
 
-from ...data import (
-    CourseLearningSequenceData, CourseOutlineData, CourseSectionData, VisibilityData, CourseVisibility
-)
+from ...data import CourseLearningSequenceData, CourseOutlineData, CourseSectionData, CourseVisibility, VisibilityData
 from ..outlines import (
     get_course_outline,
     get_user_course_outline,
diff --git a/openedx/core/djangoapps/schedules/management/commands/tests/test_send_course_update.py b/openedx/core/djangoapps/schedules/management/commands/tests/test_send_course_update.py
index abda561ab8..a7ce38f88d 100644
--- a/openedx/core/djangoapps/schedules/management/commands/tests/test_send_course_update.py
+++ b/openedx/core/djangoapps/schedules/management/commands/tests/test_send_course_update.py
@@ -12,6 +12,7 @@ from edx_ace.utils.date import serialize
 from mock import patch
 from six.moves import range
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.schedules import resolvers, tasks
 from openedx.core.djangoapps.schedules.config import COURSE_UPDATE_WAFFLE_FLAG
 from openedx.core.djangoapps.schedules.management.commands import send_course_update as nudge
@@ -21,7 +22,6 @@ from openedx.core.djangoapps.schedules.management.commands.tests.send_email_base
 )
 from openedx.core.djangoapps.schedules.management.commands.tests.upsell_base import ScheduleUpsellTestMixin
 from openedx.core.djangoapps.schedules.models import ScheduleExperience
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import skip_unless_lms
 from student.tests.factories import CourseEnrollmentFactory
 from xmodule.modulestore.tests.django_utils import ModuleStoreTestCase
diff --git a/openedx/core/djangoapps/schedules/tests/test_content_highlights.py b/openedx/core/djangoapps/schedules/tests/test_content_highlights.py
index ad96989b35..1538342ded 100644
--- a/openedx/core/djangoapps/schedules/tests/test_content_highlights.py
+++ b/openedx/core/djangoapps/schedules/tests/test_content_highlights.py
@@ -3,12 +3,14 @@
 import datetime
 from unittest.mock import patch
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.schedules.config import COURSE_UPDATE_WAFFLE_FLAG
 from openedx.core.djangoapps.schedules.content_highlights import (
-    course_has_highlights, get_week_highlights, get_next_section_highlights,
+    course_has_highlights,
+    get_next_section_highlights,
+    get_week_highlights
 )
 from openedx.core.djangoapps.schedules.exceptions import CourseUpdateDoesNotExist
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import skip_unless_lms
 from student.models import CourseEnrollment
 from student.tests.factories import UserFactory
diff --git a/openedx/core/djangoapps/schedules/tests/test_resolvers.py b/openedx/core/djangoapps/schedules/tests/test_resolvers.py
index 6ce386f27f..94ae1bcc3f 100644
--- a/openedx/core/djangoapps/schedules/tests/test_resolvers.py
+++ b/openedx/core/djangoapps/schedules/tests/test_resolvers.py
@@ -12,17 +12,17 @@ from django.test.utils import override_settings
 from testfixtures import LogCapture
 from waffle.testutils import override_switch
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.schedules.config import COURSE_UPDATE_WAFFLE_FLAG
 from openedx.core.djangoapps.schedules.models import Schedule
 from openedx.core.djangoapps.schedules.resolvers import (
     LOG,
     BinnedSchedulesBaseResolver,
-    CourseUpdateResolver,
     CourseNextSectionUpdate,
+    CourseUpdateResolver
 )
 from openedx.core.djangoapps.schedules.tests.factories import ScheduleConfigFactory
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteConfigurationFactory, SiteFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import CacheIsolationMixin, skip_unless_lms
 from student.tests.factories import CourseEnrollmentFactory
 from xmodule.modulestore.tests.django_utils import ModuleStoreTestCase
diff --git a/openedx/core/djangoapps/schedules/tests/test_signals.py b/openedx/core/djangoapps/schedules/tests/test_signals.py
index e397a51413..ea01baac59 100644
--- a/openedx/core/djangoapps/schedules/tests/test_signals.py
+++ b/openedx/core/djangoapps/schedules/tests/test_signals.py
@@ -7,23 +7,24 @@ import datetime
 
 import ddt
 import pytest
-from course_modes.models import CourseMode
-from course_modes.tests.factories import CourseModeFactory
 from mock import patch
 from pytz import utc
-from student.models import CourseEnrollment
-from student.tests.factories import CourseEnrollmentFactory
 from testfixtures import LogCapture
-from xmodule.modulestore import ModuleStoreEnum
-from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
-from xmodule.modulestore.tests.factories import CourseFactory
 
+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.courseware.models import DynamicUpgradeDeadlineConfiguration
 from openedx.core.djangoapps.schedules.models import ScheduleExperience
 from openedx.core.djangoapps.schedules.signals import CREATE_SCHEDULE_WAFFLE_FLAG, log
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import skip_unless_lms
+from student.models import CourseEnrollment
+from student.tests.factories import CourseEnrollmentFactory
+from xmodule.modulestore import ModuleStoreEnum
+from xmodule.modulestore.tests.django_utils import SharedModuleStoreTestCase
+from xmodule.modulestore.tests.factories import CourseFactory
+
 from ..models import Schedule
 from ..tests.factories import ScheduleConfigFactory
 
diff --git a/openedx/core/djangoapps/user_api/accounts/tests/test_settings_views.py b/openedx/core/djangoapps/user_api/accounts/tests/test_settings_views.py
index b235c7f4b7..f1778123cf 100644
--- a/openedx/core/djangoapps/user_api/accounts/tests/test_settings_views.py
+++ b/openedx/core/djangoapps/user_api/accounts/tests/test_settings_views.py
@@ -12,6 +12,7 @@ from django.test.utils import override_settings
 from django.urls import reverse
 from edx_rest_api_client import exceptions
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.commerce.models import CommerceConfiguration
 from lms.djangoapps.commerce.tests import factories
 from lms.djangoapps.commerce.tests.mocks import mock_get_orders
@@ -23,11 +24,10 @@ from openedx.core.djangoapps.site_configuration.tests.mixins import SiteMixin
 from openedx.core.djangoapps.user_api.accounts.settings_views import account_settings_context, get_user_orders
 from openedx.core.djangoapps.user_api.accounts.toggles import REDIRECT_TO_ACCOUNT_MICROFRONTEND
 from openedx.core.djangoapps.user_api.tests.factories import UserPreferenceFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import skip_unless_lms
+from openedx.features.enterprise_support.utils import get_enterprise_readonly_account_fields
 from student.tests.factories import UserFactory
 from third_party_auth.tests.testutil import ThirdPartyAuthTestMixin
-from openedx.features.enterprise_support.utils import get_enterprise_readonly_account_fields
 
 
 @skip_unless_lms
diff --git a/openedx/core/djangoapps/user_authn/views/tests/test_login.py b/openedx/core/djangoapps/user_authn/views/tests/test_login.py
index df89781a19..cf2157e0b7 100644
--- a/openedx/core/djangoapps/user_authn/views/tests/test_login.py
+++ b/openedx/core/djangoapps/user_authn/views/tests/test_login.py
@@ -26,12 +26,12 @@ from openedx.core.djangoapps.password_policy.compliance import (
 )
 from openedx.core.djangoapps.user_api.accounts import EMAIL_MIN_LENGTH, EMAIL_MAX_LENGTH
 from openedx.core.djangoapps.user_authn.cookies import jwt_cookies
+from openedx.core.djangoapps.user_authn.tests.utils import setup_login_oauth_client
 from openedx.core.djangoapps.user_authn.views.login import (
-    AllowedAuthUser,
     ENABLE_LOGIN_USING_THIRDPARTY_AUTH_ONLY,
+    AllowedAuthUser,
     _check_user_auth_flow
 )
-from openedx.core.djangoapps.user_authn.tests.utils import setup_login_oauth_client
 from openedx.core.djangolib.testing.utils import CacheIsolationTestCase, skip_unless_lms
 from openedx.core.djangoapps.site_configuration.tests.mixins import SiteMixin
 from openedx.core.lib.api.test_utils import ApiTestCase
diff --git a/openedx/core/djangoapps/user_authn/views/tests/test_logistration.py b/openedx/core/djangoapps/user_authn/views/tests/test_logistration.py
index 04ac77b0a2..cac9e02b2e 100644
--- a/openedx/core/djangoapps/user_authn/views/tests/test_logistration.py
+++ b/openedx/core/djangoapps/user_authn/views/tests/test_logistration.py
@@ -23,6 +23,7 @@ from pytz import UTC
 from six.moves.urllib.parse import urlencode  # pylint: disable=import-error
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.branding.api import get_privacy_url
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteConfigurationFactory, SiteFactory
 from openedx.core.djangoapps.site_configuration.tests.mixins import SiteMixin
diff --git a/openedx/core/djangoapps/user_authn/views/tests/test_register.py b/openedx/core/djangoapps/user_authn/views/tests/test_register.py
index 86b044aa21..aaf8da2133 100644
--- a/openedx/core/djangoapps/user_authn/views/tests/test_register.py
+++ b/openedx/core/djangoapps/user_authn/views/tests/test_register.py
@@ -21,6 +21,7 @@ from pytz import UTC
 from six.moves import range
 from social_django.models import Partial, UserSocialAuth
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.site_configuration.helpers import get_value
 from openedx.core.djangoapps.site_configuration.tests.test_util import with_site_configuration
 from openedx.core.djangoapps.user_api.accounts import (
@@ -49,7 +50,6 @@ from openedx.core.djangoapps.user_api.tests.test_constants import SORTED_COUNTRI
 from openedx.core.djangoapps.user_api.tests.test_helpers import TestCaseForm
 from openedx.core.djangoapps.user_api.tests.test_views import UserAPITestCase
 from openedx.core.djangoapps.user_authn.views.register import REGISTRATION_FAILURE_LOGGING_FLAG
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import CacheIsolationTestCase, skip_unless_lms
 from openedx.core.lib.api import test_utils
 from student.helpers import authenticate_new_user
diff --git a/openedx/core/djangoapps/waffle_utils/__init__.py b/openedx/core/djangoapps/waffle_utils/__init__.py
index 086d97e287..98a3b4afd7 100644
--- a/openedx/core/djangoapps/waffle_utils/__init__.py
+++ b/openedx/core/djangoapps/waffle_utils/__init__.py
@@ -7,8 +7,7 @@ from contextlib import contextmanager
 
 from opaque_keys.edx.keys import CourseKey
 
-from edx_toggles.toggles import WaffleFlag as BaseWaffleFlag
-from edx_toggles.toggles import WaffleFlagNamespace
+from edx_toggles.toggles import WaffleFlag, WaffleFlagNamespace
 from edx_toggles.toggles import WaffleSwitch as BaseWaffleSwitch
 from edx_toggles.toggles import WaffleSwitchNamespace as BaseWaffleSwitchNamespace
 
@@ -94,26 +93,6 @@ class WaffleSwitch(BaseWaffleSwitch):
             yield
 
 
-class WaffleFlag(BaseWaffleFlag):
-    """
-    Waffle flag class that implements custom override method.
-
-    This class should be removed in favour of edx_toggles.toggles.WaffleFlag once we get rid of the WaffleFlagNamespace
-    class and the `override` method.
-    """
-
-    @contextmanager
-    def override(self, active=True):
-        """
-        Shortcut method for `override_waffle_flag`.
-        """
-        # 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):
-            yield
-
 
 class CourseWaffleFlag(WaffleFlag):
     """
diff --git a/openedx/features/calendar_sync/tests/test_plugins.py b/openedx/features/calendar_sync/tests/test_plugins.py
index 077b3c87f9..52d8132c83 100644
--- a/openedx/features/calendar_sync/tests/test_plugins.py
+++ b/openedx/features/calendar_sync/tests/test_plugins.py
@@ -7,8 +7,8 @@ import crum
 import ddt
 from django.test import RequestFactory
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 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
diff --git a/openedx/features/content_type_gating/tests/test_models.py b/openedx/features/content_type_gating/tests/test_models.py
index 89c66668d0..e58aa32ac4 100644
--- a/openedx/features/content_type_gating/tests/test_models.py
+++ b/openedx/features/content_type_gating/tests/test_models.py
@@ -11,10 +11,10 @@ from mock import Mock
 from opaque_keys.edx.locator import CourseLocator
 
 from course_modes.tests.factories import CourseModeFactory
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangoapps.config_model_utils.models import Provenance
 from openedx.core.djangoapps.content.course_overviews.tests.factories import CourseOverviewFactory
 from openedx.core.djangoapps.site_configuration.tests.factories import SiteConfigurationFactory
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import CacheIsolationTestCase
 from openedx.features.content_type_gating.models import ContentTypeGatingConfig
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
diff --git a/openedx/features/course_experience/tests/views/test_course_home.py b/openedx/features/course_experience/tests/views/test_course_home.py
index ca533fdaf7..bdb35d3265 100644
--- a/openedx/features/course_experience/tests/views/test_course_home.py
+++ b/openedx/features/course_experience/tests/views/test_course_home.py
@@ -20,12 +20,11 @@ from waffle.testutils import override_flag
 
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
-from lms.djangoapps.courseware.tests.helpers import get_expiration_banner_text
+from edx_toggles.toggles.testutils import override_waffle_flag
 from experiments.models import ExperimentData
 from lms.djangoapps.commerce.models import CommerceConfiguration
 from lms.djangoapps.commerce.utils import EcommerceService
 from lms.djangoapps.course_goals.api import add_course_goal, remove_course_goal
-from lms.djangoapps.courseware.utils import verified_upgrade_deadline_link
 from lms.djangoapps.courseware.tests.factories import (
     BetaTesterFactory,
     GlobalStaffFactory,
@@ -34,9 +33,9 @@ from lms.djangoapps.courseware.tests.factories import (
     OrgStaffFactory,
     StaffFactory
 )
+from lms.djangoapps.courseware.tests.helpers import get_expiration_banner_text
+from lms.djangoapps.courseware.utils import verified_upgrade_deadline_link
 from lms.djangoapps.discussion.django_comment_client.tests.factories import RoleFactory
-from openedx.features.discounts.applicability import get_discount_expiration_date
-from openedx.features.discounts.utils import format_strikeout_price, REV1008_EXPERIMENT_ID
 from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
 from openedx.core.djangoapps.dark_lang.models import DarkLangConfig
 from openedx.core.djangoapps.django_comment_common.models import (
@@ -55,6 +54,8 @@ from openedx.features.course_experience import (
     SHOW_REVIEWS_TOOL_FLAG,
     SHOW_UPGRADE_MSG_ON_COURSE_HOME
 )
+from openedx.features.discounts.applicability import get_discount_expiration_date
+from openedx.features.discounts.utils import REV1008_EXPERIMENT_ID, format_strikeout_price
 from student.models import CourseEnrollment, FBEEnrollmentExclusion
 from student.tests.factories import UserFactory
 from util.date_utils import strftime_localized
@@ -1021,7 +1022,7 @@ class CourseHomeFragmentViewTests(ModuleStoreTestCase):
     def test_upgrade_message_discount(self):
         CourseEnrollment.enroll(self.user, self.course.id, CourseMode.AUDIT)
 
-        with SHOW_UPGRADE_MSG_ON_COURSE_HOME.override(True):
+        with override_waffle_flag(SHOW_UPGRADE_MSG_ON_COURSE_HOME, True):
             response = self.client.get(self.url)
 
         self.assertContains(response, "<span>DISCOUNT_PRICE</span>")
diff --git a/openedx/features/course_experience/tests/views/test_course_sock.py b/openedx/features/course_experience/tests/views/test_course_sock.py
index bce7a4ab2b..1ac8bd2b0a 100644
--- a/openedx/features/course_experience/tests/views/test_course_sock.py
+++ b/openedx/features/course_experience/tests/views/test_course_sock.py
@@ -3,13 +3,12 @@ Tests for course verification sock
 """
 
 
+import ddt
 import mock
 
-import ddt
-
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.commerce.models import CommerceConfiguration
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.core.djangolib.markup import HTML
 from openedx.features.course_experience import DISPLAY_COURSE_SOCK_FLAG
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
diff --git a/openedx/features/course_experience/tests/views/test_masquerade.py b/openedx/features/course_experience/tests/views/test_masquerade.py
index 9c11710b53..0072a69612 100644
--- a/openedx/features/course_experience/tests/views/test_masquerade.py
+++ b/openedx/features/course_experience/tests/views/test_masquerade.py
@@ -2,8 +2,8 @@
 Tests for masquerading functionality on course_experience
 """
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.courseware.tests.helpers import MasqueradeMixin
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.course_experience import DISPLAY_COURSE_SOCK_FLAG, SHOW_UPGRADE_MSG_ON_COURSE_HOME
 from student.roles import CourseStaffRole
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
diff --git a/openedx/features/discounts/tests/test_applicability.py b/openedx/features/discounts/tests/test_applicability.py
index b00ec636ff..716ce212cc 100644
--- a/openedx/features/discounts/tests/test_applicability.py
+++ b/openedx/features/discounts/tests/test_applicability.py
@@ -14,10 +14,10 @@ from mock import Mock, patch
 
 from course_modes.models import CourseMode
 from course_modes.tests.factories import CourseModeFactory
+from edx_toggles.toggles.testutils import override_waffle_flag
 from entitlements.tests.factories import CourseEntitlementFactory
 from experiments.models import ExperimentData
 from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.discounts.models import DiscountRestrictionConfig
 from openedx.features.discounts.utils import REV1008_EXPERIMENT_ID
 from student.tests.factories import CourseEnrollmentFactory, UserFactory
diff --git a/openedx/features/enterprise_support/tests/test_utils.py b/openedx/features/enterprise_support/tests/test_utils.py
index f6e97faf35..3c8b106f82 100644
--- a/openedx/features/enterprise_support/tests/test_utils.py
+++ b/openedx/features/enterprise_support/tests/test_utils.py
@@ -3,19 +3,20 @@ Test the enterprise support utils.
 """
 
 import json
-import mock
-import ddt
 
+import ddt
+import mock
 from django.test import TestCase
 from django.test.utils import override_settings
 
+from edx_toggles.toggles.testutils import override_waffle_flag
 from openedx.core.djangolib.testing.utils import skip_unless_lms
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
-from openedx.features.enterprise_support.utils import ENTERPRISE_HEADER_LINKS, get_enterprise_learner_portal
 from openedx.features.enterprise_support.tests import FEATURES_WITH_ENTERPRISE_ENABLED
 from openedx.features.enterprise_support.tests.factories import (
-    EnterpriseCustomerBrandingConfigurationFactory, EnterpriseCustomerUserFactory,
+    EnterpriseCustomerBrandingConfigurationFactory,
+    EnterpriseCustomerUserFactory
 )
+from openedx.features.enterprise_support.utils import ENTERPRISE_HEADER_LINKS, get_enterprise_learner_portal
 from student.tests.factories import UserFactory
 
 
diff --git a/openedx/features/learner_profile/tests/views/test_learner_profile.py b/openedx/features/learner_profile/tests/views/test_learner_profile.py
index 58e034488f..ec97ac1450 100644
--- a/openedx/features/learner_profile/tests/views/test_learner_profile.py
+++ b/openedx/features/learner_profile/tests/views/test_learner_profile.py
@@ -13,12 +13,12 @@ from django.urls import reverse
 from opaque_keys.edx.locator import CourseLocator
 
 from course_modes.models import CourseMode
+from edx_toggles.toggles.testutils import override_waffle_flag
 from lms.djangoapps.certificates.api import is_passing_status
 from lms.djangoapps.certificates.tests.factories import GeneratedCertificateFactory
 from lms.envs.test import CREDENTIALS_PUBLIC_SERVICE_URL
 from openedx.core.djangoapps.content.course_overviews.models import CourseOverview
 from openedx.core.djangoapps.site_configuration.tests.mixins import SiteMixin
-from openedx.core.djangoapps.waffle_utils.testutils import override_waffle_flag
 from openedx.features.learner_profile.toggles import REDIRECT_TO_PROFILE_MICROFRONTEND
 from openedx.features.learner_profile.views.learner_profile import learner_profile_context
 from student.tests.factories import CourseEnrollmentFactory, UserFactory