Testing python memo decorator

Question

I do have following decorator:

def memo(f):
        """Decorator that caches the return value for each call to f(args).
        Then when called again with same args, we can just look it up."""
        cache = {}

        def _f(*args):
            try:
                return cache[args]
            except KeyError:
                cache[args] = result = f(*args)
                return result
            except TypeError:
                # some element of args can't be a dict key
                return f(args)

        return _f

I need to write some tests to know if it works good. How can I test such decorator? I only managed to write test for performance to know if it speeds up functions.

def test_memo(self):
        def fib(n):
            if n == 0:
                return 0
            elif n == 1:
                return 1
            else:
                return fib(n - 1) + fib(n - 2)

        @memo
        def cached_fib(n):
            if n == 0:
                return 0
            elif n == 1:
                return 1
            else:
                return cached_fib(n - 1) + cached_fib(n - 2)

        t0 = time.clock()
        fib(20)
        t = time.clock() - t0

        t1 = time.clock()
        cached_fib(20)
        t2 = time.clock() - t1
        self.assertGreater(t, t2)

Maybe it would be wise to test if it store some value in cache but I do not know how to achieve it in python. Any ideas?

Answer 1

The function decoarted should be called once for the same argument. Check that.

def test_memo__function_should_be_called_once_for_same_arg(self):
    @memo
    def f(arg):
        f.call_count += 1
        return arg
    f.call_count = 0

    self.assertEqual(f(1), 1)
    self.assertEqual(f(1), 1)
    self.assertEqual(f.call_count, 1)

    self.assertEqual(f(2), 2)
    self.assertEqual(f(2), 2)
    self.assertEqual(f(2), 2)
    self.assertEqual(f.call_count, 2)

BTW, In the cached_fib function, it should call cache_fib, not fib to take advantage of the memoization.