From 4ccf7154c3afdf45825857cb4249eb30744c9f74 Mon Sep 17 00:00:00 2001 From: Simon Charette Date: Thu, 19 Nov 2015 22:27:37 -0500 Subject: [PATCH] Unified test context decorators. Thanks to Tim for the review. --- django/test/utils.py | 263 +++++++++++++++++++++---------------------- 1 file changed, 131 insertions(+), 132 deletions(-) diff --git a/django/test/utils.py b/django/test/utils.py index 73f30ef6ca..ee2837f497 100644 --- a/django/test/utils.py +++ b/django/test/utils.py @@ -5,7 +5,7 @@ import time import warnings from contextlib import contextmanager from functools import wraps -from unittest import skipIf, skipUnless +from unittest import TestCase, skipIf, skipUnless from xml.dom.minidom import Node, parseString from django.apps import apps @@ -20,7 +20,7 @@ from django.template import Template from django.test.signals import setting_changed, template_rendered from django.urls import get_script_prefix, set_script_prefix from django.utils import six -from django.utils.decorators import ContextDecorator +from django.utils.decorators import available_attrs from django.utils.encoding import force_str from django.utils.translation import deactivate @@ -151,45 +151,81 @@ def get_runner(settings, test_runner_class=None): return test_runner -class override_settings(object): +class TestContextDecorator(object): """ - Acts as either a decorator, or a context manager. If it's a decorator it + A base class that can either be used as a context manager during tests + or as a test function or unittest.TestCase subclass decorator to perform + temporary alterations. + + `attr_name`: attribute assigned the return value of enable() if used as + a class decorator. + + `kwarg_name`: keyword argument passing the return value of enable() if + used as a function decorator. + """ + def __init__(self, attr_name=None, kwarg_name=None): + self.attr_name = attr_name + self.kwarg_name = kwarg_name + + def enable(self): + raise NotImplementedError + + def disable(self): + raise NotImplementedError + + def __enter__(self): + return self.enable() + + def __exit__(self, exc_type, exc_value, traceback): + self.disable() + + def decorate_class(self, cls): + if issubclass(cls, TestCase): + decorated_setUp = cls.setUp + decorated_tearDown = cls.tearDown + + def setUp(inner_self): + context = self.enable() + if self.attr_name: + setattr(inner_self, self.attr_name, context) + decorated_setUp(inner_self) + + def tearDown(inner_self): + decorated_tearDown(inner_self) + self.disable() + + cls.setUp = setUp + cls.tearDown = tearDown + return cls + raise TypeError('Can only decorate subclasses of unittest.TestCase') + + def decorate_callable(self, func): + @wraps(func, assigned=available_attrs(func)) + def inner(*args, **kwargs): + with self as context: + if self.kwarg_name: + kwargs[self.kwarg_name] = context + return func(*args, **kwargs) + return inner + + def __call__(self, decorated): + if isinstance(decorated, type): + return self.decorate_class(decorated) + elif callable(decorated): + return self.decorate_callable(decorated) + raise TypeError('Cannot decorate object of type %s' % type(decorated)) + + +class override_settings(TestContextDecorator): + """ + Acts as either a decorator or a context manager. If it's a decorator it takes a function and returns a wrapped function. If it's a contextmanager it's used with the ``with`` statement. In either event entering/exiting are called before and after, respectively, the function/block is executed. """ def __init__(self, **kwargs): self.options = kwargs - - def __enter__(self): - self.enable() - - def __exit__(self, exc_type, exc_value, traceback): - self.disable() - - def __call__(self, test_func): - from django.test import SimpleTestCase - if isinstance(test_func, type): - if not issubclass(test_func, SimpleTestCase): - raise Exception( - "Only subclasses of Django SimpleTestCase can be decorated " - "with override_settings") - self.save_options(test_func) - return test_func - else: - @wraps(test_func) - def inner(*args, **kwargs): - with self: - return test_func(*args, **kwargs) - return inner - - def save_options(self, test_func): - if test_func._overridden_settings is None: - test_func._overridden_settings = self.options - else: - # Duplicate dict to prevent subclasses from altering their parent. - test_func._overridden_settings = dict( - test_func._overridden_settings, **self.options) + super(override_settings, self).__init__() def enable(self): # Keep this code at the beginning to leave the settings unchanged @@ -219,6 +255,23 @@ class override_settings(object): setting_changed.send(sender=settings._wrapped.__class__, setting=key, value=new_value, enter=False) + def save_options(self, test_func): + if test_func._overridden_settings is None: + test_func._overridden_settings = self.options + else: + # Duplicate dict to prevent subclasses from altering their parent. + test_func._overridden_settings = dict( + test_func._overridden_settings, **self.options) + + def decorate_class(self, cls): + from django.test import SimpleTestCase + if not issubclass(cls, SimpleTestCase): + raise ValueError( + "Only subclasses of Django SimpleTestCase can be decorated " + "with override_settings") + self.save_options(cls) + return cls + class modify_settings(override_settings): """ @@ -233,6 +286,7 @@ class modify_settings(override_settings): else: assert not args self.operations = list(kwargs.items()) + super(override_settings, self).__init__() def save_options(self, test_func): if test_func._modified_settings is None: @@ -267,28 +321,29 @@ class modify_settings(override_settings): super(modify_settings, self).enable() -def override_system_checks(new_checks, deployment_checks=None): - """ Acts as a decorator. Overrides list of registered system checks. +class override_system_checks(TestContextDecorator): + """ + Acts as a decorator. Overrides list of registered system checks. Useful when you override `INSTALLED_APPS`, e.g. if you exclude `auth` app, - you also need to exclude its system checks. """ + you also need to exclude its system checks. + """ + def __init__(self, new_checks, deployment_checks=None): + from django.core.checks.registry import registry + self.registry = registry + self.new_checks = new_checks + self.deployment_checks = deployment_checks + super(override_system_checks, self).__init__() - from django.core.checks.registry import registry + def enable(self): + self.old_checks = self.registry.registered_checks + self.registry.registered_checks = self.new_checks + self.old_deployment_checks = self.registry.deployment_checks + if self.deployment_checks is not None: + self.registry.deployment_checks = self.deployment_checks - def outer(test_func): - @wraps(test_func) - def inner(*args, **kwargs): - old_checks = registry.registered_checks - registry.registered_checks = new_checks - old_deployment_checks = registry.deployment_checks - if deployment_checks is not None: - registry.deployment_checks = deployment_checks - try: - return test_func(*args, **kwargs) - finally: - registry.registered_checks = old_checks - registry.deployment_checks = old_deployment_checks - return inner - return outer + def disable(self): + self.registry.registered_checks = self.old_checks + self.registry.deployment_checks = self.old_deployment_checks def compare_xml(want, got): @@ -428,40 +483,22 @@ class CaptureQueriesContext(object): self.final_queries = len(self.connection.queries_log) -class ignore_warnings(object): +class ignore_warnings(TestContextDecorator): def __init__(self, **kwargs): self.ignore_kwargs = kwargs if 'message' in self.ignore_kwargs or 'module' in self.ignore_kwargs: self.filter_func = warnings.filterwarnings else: self.filter_func = warnings.simplefilter + super(ignore_warnings, self).__init__() - def __call__(self, decorated): - if isinstance(decorated, type): - # A class is decorated - saved_setUp = decorated.setUp - saved_tearDown = decorated.tearDown + def enable(self): + self.catch_warnings = warnings.catch_warnings() + self.catch_warnings.__enter__() + self.filter_func('ignore', **self.ignore_kwargs) - def setUp(inner_self): - self.catch_warnings = warnings.catch_warnings() - self.catch_warnings.__enter__() - self.filter_func('ignore', **self.ignore_kwargs) - saved_setUp(inner_self) - - def tearDown(inner_self): - saved_tearDown(inner_self) - self.catch_warnings.__exit__(*sys.exc_info()) - - decorated.setUp = setUp - decorated.tearDown = tearDown - return decorated - else: - @wraps(decorated) - def inner(*args, **kwargs): - with warnings.catch_warnings(): - self.filter_func('ignore', **self.ignore_kwargs) - return decorated(*args, **kwargs) - return inner + def disable(self): + self.catch_warnings.__exit__(*sys.exc_info()) @contextmanager @@ -610,23 +647,20 @@ def require_jinja2(test_func): return test_func -class ScriptPrefix(ContextDecorator): - def __enter__(self): - set_script_prefix(self.prefix) - - def __exit__(self, exc_type, exc_val, traceback): - set_script_prefix(self.old_prefix) - - def __init__(self, prefix): - self.prefix = prefix - self.old_prefix = get_script_prefix() - - -def override_script_prefix(prefix): +class override_script_prefix(TestContextDecorator): """ Decorator or context manager to temporary override the script prefix. """ - return ScriptPrefix(prefix) + def __init__(self, prefix): + self.prefix = prefix + super(override_script_prefix, self).__init__() + + def enable(self): + self.old_prefix = get_script_prefix() + set_script_prefix(self.prefix) + + def disable(self): + set_script_prefix(self.old_prefix) class LoggingCaptureMixin(object): @@ -644,7 +678,7 @@ class LoggingCaptureMixin(object): self.logger.handlers[0].stream = self.old_stream -class isolate_apps(object): +class isolate_apps(TestContextDecorator): """ Act as either a decorator or a context manager to register models defined in its wrapped context to an isolated registry. @@ -657,14 +691,13 @@ class isolate_apps(object): `attr_name`: attribute assigned the isolated registry if used as a class decorator. - `kwarg_name`: keyword argument passing the isolated registry to the - decorated method. + `kwarg_name`: keyword argument passing the isolated registry if used as a + function decorator. """ def __init__(self, *installed_apps, **kwargs): self.installed_apps = installed_apps - self.attr_name = kwargs.pop('attr_name', None) - self.kwarg_name = kwargs.pop('kwarg_name', None) + super(isolate_apps, self).__init__(**kwargs) def enable(self): self.old_apps = Options.default_apps @@ -674,37 +707,3 @@ class isolate_apps(object): def disable(self): setattr(Options, 'default_apps', self.old_apps) - - def __enter__(self): - return self.enable() - - def __exit__(self, exc_type, exc_value, traceback): - self.disable() - - def __call__(self, decorated): - if isinstance(decorated, type): - # A class is decorated - decorated_setUp = decorated.setUp - decorated_tearDown = decorated.tearDown - - def setUp(inner_self): - apps = self.enable() - if self.attr_name: - setattr(inner_self, self.attr_name, apps) - decorated_setUp(inner_self) - - def tearDown(inner_self): - decorated_tearDown(inner_self) - self.disable() - - decorated.setUp = setUp - decorated.tearDown = tearDown - return decorated - else: - @wraps(decorated) - def inner(*args, **kwargs): - with self as apps: - if self.kwarg_name: - kwargs[self.kwarg_name] = apps - return decorated(*args, **kwargs) - return inner