[Culprit-Finder] Add generator cache decorator.

appengine/findit/libs/test/cache_decorator_test.py

Index: appengine/findit/libs/test/cache_decorator_test.py
diff --git a/appengine/findit/libs/test/cache_decorator_test.py b/appengine/findit/libs/test/cache_decorator_test.py
index 03e31d976bb7cc943625a0de13f9a894993a35f4..2f348ecd4dd49ca14aa0d7b037e8a571c22901e8 100644
--- a/appengine/findit/libs/test/cache_decorator_test.py
+++ b/appengine/findit/libs/test/cache_decorator_test.py
@@ -2,11 +2,11 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import mock
 import pickle
+import unittest
 import zlib
 
-from testing_utils import testing
-
 from libs import cache
 from libs import cache_decorator
 
@@ -22,25 +22,30 @@ class _DummyCache(cache.Cache):
     self.cached_data[key] = data
 
 
-def _DummyKeyGenerator(func, *_):
-  return func.__name__
+def _DummyKeyGenerator(func, args, kwargs,  # pylint: disable=W0613
+                       namespace=None):
+  namespace = namespace or '%s.%s' % (func.__module__, func.__name__)
+  return namespace + '-' + func.__name__
+
 
+class CachedTest(unittest.TestCase):
+  """Tests ``Cached`` decorator."""
 
-class CacheDecoratorTest(testing.AppengineTestCase):
   def testDefaultKeyGenerator(self):
+    namespace = 'n'
     expected_params = {
         'id1': 'fi',
         'id2': 'pi',
         'url': 'http://url',
     }
     # Hexadecimal digits of MD5 digest of "pickled_params".
-    expected_key = 'f5f173c811f7c537a80d44511903a3e0'
+    expected_key = namespace + '-f5f173c811f7c537a80d44511903a3e0'
 
     def MockPickleDumps(params):
       self.assertEqual(expected_params, params)
       return 'pickled_params'
 
-    def Func(id1, id2, url=None):  # Unused parameters-pylint: disable=W0613
+    def func(id1, id2, url=None):  # Unused parameters-pylint: disable=W0613
       return 1  # pragma: no cover.
 
     class CallableIdentifier(object):
@@ -52,34 +57,34 @@ class CacheDecoratorTest(testing.AppengineTestCase):
       def identifier(self):
         return 'pi'
 
-    self.mock(pickle, 'dumps', MockPickleDumps)
-
-    args = (CallableIdentifier(), PropertyIdentifier())
-    kwargs = {'url': 'http://url'}
-    key = cache_decorator._DefaultKeyGenerator(Func, args, kwargs)
-    self.assertEqual(expected_key, key)
+    with mock.patch('pickle.dumps', MockPickleDumps):
+      args = (CallableIdentifier(), PropertyIdentifier())
+      kwargs = {'url': 'http://url'}
+      key = cache_decorator._DefaultKeyGenerator(
+          func, args, kwargs, namespace=namespace)
+      self.assertEqual(expected_key, key)
 
   def testCachedDecoratorWhenResultIsAlreadyCached(self):
-    dummy_cache = _DummyCache({'n-Func': 1})
+    dummy_cache = _DummyCache({'n-func': 1})
 
     @cache_decorator.Cached(
         namespace='n', key_generator=_DummyKeyGenerator, cache=dummy_cache)
-    def Func():
+    def func():
       return 2  # pragma: no cover.
 
-    self.assertEqual(1, Func())
-    self.assertEqual({'n-Func': 1}, dummy_cache.cached_data)
+    self.assertEqual(1, func())
+    self.assertEqual({'n-func': 1}, dummy_cache.cached_data)
 
   def testCachedDecoratorWhenResultIsNotCachedYet(self):
     dummy_cache = _DummyCache({})
 
     @cache_decorator.Cached(
         namespace='n', key_generator=_DummyKeyGenerator, cache=dummy_cache)
-    def Func():
+    def func():
       return 2
 
-    self.assertEqual(2, Func())
-    self.assertEqual({'n-Func': 2}, dummy_cache.cached_data)
+    self.assertEqual(2, func())
+    self.assertEqual({'n-func': 2}, dummy_cache.cached_data)
 
   def testCachedDecoratorWhenResultShouldNotBeCached(self):
     dummy_cache = _DummyCache({})
@@ -88,18 +93,18 @@ class CacheDecoratorTest(testing.AppengineTestCase):
 
     @cache_decorator.Cached(
         namespace='n', key_generator=_DummyKeyGenerator, cache=dummy_cache)
-    def Func():
+    def func():
       return results.pop()
 
-    self.assertEqual('', Func())
+    self.assertEqual('', func())
     self.assertEqual({}, dummy_cache.cached_data)
-    self.assertEqual({}, Func())
+    self.assertEqual({}, func())
     self.assertEqual({}, dummy_cache.cached_data)
-    self.assertEqual([], Func())
+    self.assertEqual([], func())
     self.assertEqual({}, dummy_cache.cached_data)
-    self.assertEqual(0, Func())
+    self.assertEqual(0, func())
     self.assertEqual({}, dummy_cache.cached_data)
-    self.assertIsNone(Func())
+    self.assertIsNone(func())
     self.assertEqual({}, dummy_cache.cached_data)
 
   def testCachedDecoratorWithMethodInAClass(self):
@@ -114,15 +119,78 @@ class CacheDecoratorTest(testing.AppengineTestCase):
         return self.url
 
       @cache_decorator.Cached(cache=_DummyCache({}))
-      def Func(self, path):
+      def func(self, path):
         self.runs += 1
         return self.url + '/' + path
 
     a1 = A('http://test', 3)
-    self.assertEqual('http://test/p1', a1.Func('p1'))
-    self.assertEqual('http://test/p1', a1.Func('p1'))
+    self.assertEqual('http://test/p1', a1.func('p1'))
+    self.assertEqual('http://test/p1', a1.func('p1'))
     self.assertEqual(1, a1.runs)
 
     a2 = A('http://test', 5)
-    self.assertEqual('http://test/p1', a2.Func('p1'))
+    self.assertEqual('http://test/p1', a2.func('p1'))
     self.assertEqual(0, a2.runs)
+
+
+class GeneratorCachedTest(unittest.TestCase):
+  """Tests ``GeneratorCached`` decorator."""
+
+  def testWhenResultIsAlreadyCached(self):
+    value_list = [0, 1, 2]
+    key = 'n-func'
+    # Zero value won't be cached.
+    cached_keys = ['%s-%d' % (key, i) for i in xrange(len(value_list))]
+    cached_data = {key: value for key, value in zip(cached_keys, value_list)}
+    cached_data[key] = cached_keys
+
+    dummy_cache = _DummyCache(cached_data)
+
+    @cache_decorator.GeneratorCached(
+        dummy_cache, namespace='n', key_generator=_DummyKeyGenerator)
+    def func():  # pragma: no cover
+      for value in value_list:
+        yield value
+
+    for value, expected_value in zip(func(), value_list):
+      self.assertEqual(value, expected_value)
+
+  def testWhenResultIsNotCachedYet(self):
+    value_list = [0, 1, 2, 3, 4]
+    dummy_cache = _DummyCache({})
+
+    @cache_decorator.GeneratorCached(
+        dummy_cache, namespace='n', key_generator=_DummyKeyGenerator)
+    def func():  # pragma: no cover
+      for value in value_list:
+        yield value
+
+    key = _DummyKeyGenerator(func, [], {}, namespace='n')
+    cached_keys = ['%s-%d' % (key, i) for i in xrange(len(value_list))]
+    cached_data = {key: value for key, value in zip(cached_keys, value_list)}
+    cached_data[key] = cached_keys
+    for value, expected_value in zip(func(), value_list):
+      self.assertEqual(value, expected_value)
+
+    self.assertDictEqual(dummy_cache.cached_data, cached_data)
+
+  @mock.patch('libs.cache_decorator.GeneratorCached.SetCache',
+              lambda *_: False)
+  def testFailedToSetValue(self):
+    value_list = [0, 1, 2, 3, 4]
+    dummy_cache = _DummyCache({})
+
+    @cache_decorator.GeneratorCached(
+        dummy_cache, namespace='n', key_generator=_DummyKeyGenerator)
+    def func():  # pragma: no cover
+      for value in value_list:
+        yield value
+
+    key = _DummyKeyGenerator(func, [], {}, namespace='n')
+    cached_keys = ['%s-%d' % (key, i) for i in xrange(len(value_list))]
+    cached_data = {key: value for key, value in zip(cached_keys, value_list)}
+    cached_data[key] = cached_keys
+    for value, expected_value in zip(func(), value_list):
+      self.assertEqual(value, expected_value)
+
+    self.assertDictEqual(dummy_cache.cached_data, {})