Ticket #20536: test_filebasedcache.py

# coding=utf-8
import os
import tempfile
import time
from multiprocessing import Pool
from django.conf import settings
try:
    from django.utils.six.moves import cPickle as pickle
except ImportError:
    import pickle

settings.configure()

from django.core.cache.backends.filebased import FileBasedCache


class SafeFileBasedCache(FileBasedCache):
    """
    FileBasedCache backend that uses temp file write-troughs to prevent
    concurrency issues.
    """
    _tmp_suffix = '.__django_cache'

    def set(self, key, value, timeout=None, version=None):
        key = self.make_key(key, version=version)
        self.validate_key(key)

        fname = self._key_to_file(key)
        dirname = os.path.dirname(fname)

        if timeout is None:
            timeout = self.default_timeout

        self._cull()

        try:
            if not os.path.exists(dirname):
                os.makedirs(dirname)

            fd, tmp = tempfile.mkstemp(suffix=self._tmp_suffix, dir=dirname)
            with os.fdopen(fd, 'wb') as f:
                now = time.time()
                pickle.dump(now + timeout, f, pickle.HIGHEST_PROTOCOL)
                pickle.dump(value, f, pickle.HIGHEST_PROTOCOL)
            os.rename(tmp, fname)
        except (IOError, OSError):
            pass


cache = FileBasedCache('./tmp', {})
safecache = SafeFileBasedCache('./tmp', {})
text = u'Iñtërnâtiônàlizætiøn' * 1024


def test(cache):
    value = {'list1': range(1024), 'text': text, 'list': range(1024)}
    try:
        cache.set('test', value)
        cached = cache.get('test')
        if cached is None:
            return 'miss'
        elif cached != value:
            return 'fail'
    except Exception, e:
        return str(e)


def test_cache(i):
    return test(cache)


def test_safe_cache(i):
    return test(safecache)


def test_multiprocess(func):
    pool = Pool(processes=6)
    result = pool.map_async(func, xrange(512))
    pool.close()
    pool.join()
    results = result.get()
    misses = [fail for fail in results if fail == 'miss']
    fails = [fail for fail in results if fail is not None and fail != 'miss']
    print '%s miss(es), %s fail(s)' % (len(misses), len(fails))


if __name__ == '__main__':
    test_multiprocess(test_cache)
    test_multiprocess(test_safe_cache)