Updated override_waffle_flag to work as a context manager

override_waffle_flag can now be used as a context manager in addition to its previous role as a decorator. Additionally, the tests have been updated to use standard assertions since we now use py.test.

openedx/core/djangoapps/waffle_utils/tests/test_testutils.py

@@ -3,20 +3,18 @@ Tests for waffle utils test utilities.
 """
 
 import crum
-
 from django.test import TestCase
 from django.test.client import RequestFactory
 from opaque_keys.edx.keys import CourseKey
 
 from request_cache.middleware import RequestCache
-
 from .. import CourseWaffleFlag, WaffleFlagNamespace
 from ..testutils import override_waffle_flag
 
 
 class OverrideWaffleFlagTests(TestCase):
     """
-    Tests for the override_waffle_flag decorator.
+    Tests for the override_waffle_flag decorator/context manager.
     """
 
     NAMESPACE_NAME = "test_namespace"
@@ -34,30 +32,35 @@ class OverrideWaffleFlagTests(TestCase):
         RequestCache.clear_request_cache()
 
     @override_waffle_flag(TEST_COURSE_FLAG, True)
-    def check_is_enabled_with_decorator(self):
-        # test flag while overridden with decorator
-        self.assertTrue(self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY))
+    def assert_decorator_activates_flag(self):
+        assert self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY)
 
     def test_override_waffle_flag_pre_cached(self):
         # checks and caches the is_enabled value
-        self.assertFalse(self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY))
+        assert not self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY)
         flag_cache = self.TEST_COURSE_FLAG.waffle_namespace._cached_flags
-        self.assertIn(self.NAMESPACED_FLAG_NAME, flag_cache)
+        assert self.NAMESPACED_FLAG_NAME in flag_cache
 
-        # test flag while overridden with decorator
-        self.check_is_enabled_with_decorator()
+        self.assert_decorator_activates_flag()
 
         # test cached flag is restored
-        self.assertIn(self.NAMESPACED_FLAG_NAME, flag_cache)
-        self.assertEquals(self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY), False)
+        assert self.NAMESPACED_FLAG_NAME in flag_cache
+        assert not self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY)
 
     def test_override_waffle_flag_not_pre_cached(self):
         # check that the flag is not yet cached
         flag_cache = self.TEST_COURSE_FLAG.waffle_namespace._cached_flags
-        self.assertNotIn(self.NAMESPACED_FLAG_NAME, flag_cache)
+        assert self.NAMESPACED_FLAG_NAME not in flag_cache
 
-        # test flag while overridden with decorator
-        self.check_is_enabled_with_decorator()
+        self.assert_decorator_activates_flag()
 
         # test cache is removed when no longer using decorator/context manager
-        self.assertNotIn(self.NAMESPACED_FLAG_NAME, flag_cache)
+        assert self.NAMESPACED_FLAG_NAME not in flag_cache
+
+    def test_override_waffle_flag_as_context_manager(self):
+        assert not self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY)
+
+        with override_waffle_flag(self.TEST_COURSE_FLAG, True):
+            assert self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY)
+
+        assert not self.TEST_COURSE_FLAG.is_enabled(self.TEST_COURSE_KEY)

openedx/core/djangoapps/waffle_utils/testutils.py

@@ -2,8 +2,6 @@
 Test utilities for waffle utilities.
 """
 
-from functools import wraps
-
 from waffle.testutils import override_flag
 
 # Can be used with FilteredQueryCountMixin.assertNumQueries() to blacklist
@@ -13,54 +11,56 @@ from waffle.testutils import override_flag
 WAFFLE_TABLES = ['waffle_utils_waffleflagcourseoverridemodel', 'waffle_flag', 'waffle_switch', 'waffle_sample']
 
 
-def override_waffle_flag(flag, active):
+class override_waffle_flag(override_flag):
     """
-    To be used as a decorator for a test function to override a namespaced
-    waffle flag.
+    override_waffle_flag is a contextmanager for easier testing of flags.
 
-        flag (WaffleFlag): The namespaced cached waffle flag.
-        active (Boolean): The value to which the flag will be set.
+    It accepts two parameters, the flag itself and its intended state. Example
+    usage::
 
-    Example usage:
+        with override_waffle_flag(UNIFIED_COURSE_TAB_FLAG, active=True):
+            ...
 
-        @override_waffle_flag(UNIFIED_COURSE_TAB_FLAG, active=True)
+    If the flag already exists, its value will be changed inside the context
+    block, then restored to the original value. If the flag does not exist
+    before entering the context, it is created, then removed at the end of the
+    block.
 
-    """
+    It can also act as a decorator::
 
-    def real_decorator(function):
-        """
-        Actual decorator function.
+        @override_waffle_flag(UNIFIED_COURSE_TAB_FLAG, active=True)
+        def test_happy_mode_enabled():
+            ...
     """
+    _cached_value = None
 
-        @wraps(function)
-        def wrapper(*args, **kwargs):
+    def __init__(self, flag, active):
         """
-            Provides the actual override functionality of the decorator.
-
-            Saves the previous cached value of the flag and restores it (if it
-            was set), after overriding it.
 
+        Args:
+             flag (WaffleFlag): The namespaced cached waffle flag.
+             active (Boolean): The value to which the flag will be set.
         """
+        self.flag = flag
         waffle_namespace = flag.waffle_namespace
-            namespaced_flag_name = waffle_namespace._namespaced_name(flag.flag_name)
+        name = waffle_namespace._namespaced_name(flag.flag_name)  # pylint: disable=protected-access
+        super(override_waffle_flag, self).__init__(name, active)
 
-            # save previous value and whether it existed in the cache
-            cached_value_existed = namespaced_flag_name in waffle_namespace._cached_flags
-            if cached_value_existed:
-                previous_value = waffle_namespace._cached_flags[namespaced_flag_name]
+    def __enter__(self):
+        super(override_waffle_flag, self).__enter__()
 
-            # set new value
-            waffle_namespace._cached_flags[namespaced_flag_name] = active
+        # pylint: disable=protected-access
+        # Store values that have been cached on the flag
+        self._cached_value = self.flag.waffle_namespace._cached_flags.get(self.name)
+        self.flag.waffle_namespace._cached_flags[self.name] = self.active
 
-            with override_flag(namespaced_flag_name, active):
-                # call wrapped function
-                function(*args, **kwargs)
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        super(override_waffle_flag, self).__exit__(exc_type, exc_val, exc_tb)
 
-            # restore value
-            if cached_value_existed:
-                waffle_namespace._cached_flags[namespaced_flag_name] = previous_value
-            elif namespaced_flag_name in waffle_namespace._cached_flags:
-                del waffle_namespace._cached_flags[namespaced_flag_name]
-        return wrapper
+        # pylint: disable=protected-access
+        # Restore the cached values
+        waffle_namespace = self.flag.waffle_namespace
+        waffle_namespace._cached_flags.pop(self.name, None)
 
-    return real_decorator
+        if self._cached_value is not None:
+            waffle_namespace._cached_flags[self.name] = self._cached_value