From 9782b407516998e79191cafd8145f5336bf54ec3 Mon Sep 17 00:00:00 2001 From: Clinton Blackburn Date: Fri, 3 Nov 2017 15:56:26 -0400 Subject: [PATCH] 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. --- .../waffle_utils/tests/test_testutils.py | 35 +++++---- .../core/djangoapps/waffle_utils/testutils.py | 78 +++++++++---------- 2 files changed, 58 insertions(+), 55 deletions(-) diff --git a/openedx/core/djangoapps/waffle_utils/tests/test_testutils.py b/openedx/core/djangoapps/waffle_utils/tests/test_testutils.py index 107278a1f1..fa1ce59c4a 100644 --- a/openedx/core/djangoapps/waffle_utils/tests/test_testutils.py +++ b/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) diff --git a/openedx/core/djangoapps/waffle_utils/testutils.py b/openedx/core/djangoapps/waffle_utils/testutils.py index 780147fde6..b27de92c3c 100644 --- a/openedx/core/djangoapps/waffle_utils/testutils.py +++ b/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): + ... + + 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:: @override_waffle_flag(UNIFIED_COURSE_TAB_FLAG, active=True) - + def test_happy_mode_enabled(): + ... """ + _cached_value = None - def real_decorator(function): - """ - Actual decorator function. + def __init__(self, flag, active): """ - @wraps(function) - def wrapper(*args, **kwargs): - """ - Provides the actual override functionality of the decorator. + 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 + name = waffle_namespace._namespaced_name(flag.flag_name) # pylint: disable=protected-access + super(override_waffle_flag, self).__init__(name, active) - Saves the previous cached value of the flag and restores it (if it - was set), after overriding it. + def __enter__(self): + super(override_waffle_flag, self).__enter__() - """ - waffle_namespace = flag.waffle_namespace - namespaced_flag_name = waffle_namespace._namespaced_name(flag.flag_name) + # 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 - # 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 __exit__(self, exc_type, exc_val, exc_tb): + super(override_waffle_flag, self).__exit__(exc_type, exc_val, exc_tb) - # set new value - waffle_namespace._cached_flags[namespaced_flag_name] = active + # pylint: disable=protected-access + # Restore the cached values + waffle_namespace = self.flag.waffle_namespace + waffle_namespace._cached_flags.pop(self.name, None) - with override_flag(namespaced_flag_name, active): - # call wrapped function - function(*args, **kwargs) - - # 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 - - return real_decorator + if self._cached_value is not None: + waffle_namespace._cached_flags[self.name] = self._cached_value