Cache backend switching for Django tests

cache_switch.py

from django.core.cache.backends.base import BaseCache
from functools import wraps
import sys


class OldStyleClass:
    pass
ClassType = type(OldStyleClass)


inPy3k = sys.version_info[0] == 3

if inPy3k:
    class_types = type
else:
    class_types = (type, ClassType)


class CacheSwitch(BaseCache):
    def __init__(self, _location, params):
        BaseCache.__init__(self, params)
        self._current_cache = None

    def get_current_cache(self):
        if not self._current_cache:
            from django.core import cache
            self._current_cache = cache.get_cache('dummy://')
        return self._current_cache

    def set_current_cache(self, backend, **kwargs):
        if isinstance(backend, BaseCache):
            self._current_cache = backend
        else:
            from django.core import cache
            self._current_cache = cache.get_cache(backend, **kwargs)
            self._current_cache.clear()

    def get(self, key, default=None, version=None):
        cache = self.get_current_cache()
        return cache.get(key, default=default, version=version)

    def set(self, key, value, timeout=None, version=None):
        cache = self.get_current_cache()
        return cache.set(key, value, timeout=timeout, version=version)


class cache_switch(object):
    """Changes the cache backend temporarily

    This works as
    * a context manager
    * a function decorator
    * a method decorator
    * a class decorator

    Examples:

    >>> with cache_switch('locmem://'):
    ...     run_my_test()

    >>> @cache_switch('locmem://')
    ... def test_my_feature():
    ...     run_my_test()

    >>> class My_Tests:
    ...     @cache_switch('locmem://')
    ...     def test_my_feature(self):
    ...         run_my_test()

    >>> @cache_switch('locmem://')
    ... class My_Tests:
    ...     def test_my_feature(self):
    ...         run_my_test()
    """
    def __init__(self, backend, **kwargs):
        self.backend = backend
        self.kwargs = kwargs

    def __call__(self, f):
        if isinstance(f, class_types):
            return self.decorate_class(f)

        @wraps(f)
        def _inner(*args, **kw):
            self._cache_switch()
            try:
                return f(*args, **kw)
            finally:
                self._cache_unswitch()

        return _inner

    def decorate_class(self, klass):
        for attr in dir(klass):
            attr_value = getattr(klass, attr)
            if attr.startswith("test") and hasattr(attr_value, "__call__"):
                decorator = cache_switch(self.backend, **self.kwargs)
                decorated = decorator(attr_value)
                setattr(klass, attr, decorated)
        return klass

    def __enter__(self):
        self._cache_switch()

    def _get_cache_switch_backend(self):
        from django.core import cache
        if not isinstance(cache.cache, CacheSwitch):
            raise TypeError("Can't change effective cache backend "
                            "since CacheSwitch is not the default backend")
        return cache.cache

    def _cache_switch(self):
        switch = self._get_cache_switch_backend()
        self._old_cache = switch.get_current_cache()
        switch.set_current_cache(self.backend, **self.kwargs)

    def _cache_unswitch(self):
        switch = self._get_cache_switch_backend()
        switch.set_current_cache(self._old_cache)

    def __exit__(self, *args):
        self._cache_unswitch()
        return False

    start = __enter__
    stop = __exit__