TestCase base class to easily temporarily change module values

from django.conf import settings
from django.test import TestCase
from django.utils.importlib import import_module
 
 
def restore_attrs(wrapper):
    """
    Restore all attributes for a wrapped module to their original values
    (removing any which weren't originally set).
 
    """
    for key, value in wrapper._attrs_to_restore.iteritems():
        setattr(wrapper._wrapped_module, key, value)
    for setting in wrapper._delete_settings:
        try:
            delattr(wrapper._wrapped_module, setting)
        except AttributeError:
            pass
 
 
class ModuleWrapper(object):
    """
    A wrapper for modules which will remember attribute changes made so they
    can be reverted.
 
    """
 
    def __init__(self, module):
        """
        Initialise some containers to remember overwritten attributes to
        restore and newly added attributes to delete.
 
        """
        self.__dict__['_attrs_to_restore'] = {}
        self.__dict__['_attrs_to_delete'] = set()
        self.__dict__['_wrapped_module'] = module
 
    def __getattr__(self, key):
        """
        Retrieve an attribute.
 
        """
        getattr(self._wrapped_module, key)
 
    def __setattr__(self, key, value):
        """
        Change an attribute, remembering its original value so it can be
        reverted (or deleted if it didn't exist).
 
        """
        attrs_to_restore = self._attrs_to_restore
        try:
            if key not in attrs_to_restore:
                attrs_to_restore[key] = getattr(value, key)
        except AttributeError:
            self._attrs_to_delete.add(key)
        setattr(self._wrapped_module, key, value)
 
    def __delattr__(self, key):
        """
        Remove an attribute, remembering its original value so it can be
        restored.
        
        Deleting attributes which don't exist will *not* raise an exception.
 
        """
        try:
            if key not in self._attrs_to_delete and\
                            key not in self._attrs_to_restore:
                self._attrs_to_restore[key] = getattr(self._wrapped_module,
                                                      key)
            delattr(self._wrapped_module, key)
        except AttributeError:
            pass
 
 
class ModuleWrapperTestCase(TestCase):
    """
    A base test case that can be used to modify attributes in any number of
    modules under the knowledge that they will be returned to their original
    value (or removed if they didn't previously exist).
 
    To set which modules should be used, set the modules attribute on your test
    case subclass to a dictionary for modules which you would like to
    temporarily change attributes for.
    
    Each key should be the attribute which the wrapper will be available as and
    each value a module (or a dotted string representation of a module).
    
    For example::
 
        class MyTestCase(ModuleWrapperTestCase):
            modules = {'mymodule': myproject.mymodule}
            #...
 
    """
    modules = None
 
    def _pre_setup(self, *args, **kwargs):
        super(ModuleWrapperTestCase, self)._pre_setup(*args, **kwargs)
        if self.modules:
            for attr, module in self.modules.iteritems():
                if isinstance(module, basestring):
                    module = import_module(module)
                setattr(self, attr, ModuleWrapper(module))
 
    def _post_teardown(self, *args, **kwargs):
        if self.modules:
            for attr in self.modules:
                restore_attrs(getattr(self, attr))
        super(ModuleWrapperTestCase, self)._post_teardown(*args, **kwargs)
 
 
class SettingsTestCase(ModuleWrapperTestCase):
    """
    A base test case that can be used to modify settings in
    ``django.conf.settings`` under the knowledge that they will be returned
    to their original value (or removed if they didn't previously exist).
 
    """
    modules = {'settings': settings}
 
    def restore_settings(self):
        """
        Restore all settings to their original values (removing any which
        weren't originally in settings).
 
        """
        restore_attrs(self.settings)