[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..271f63905355379fbebc7c480a1a27c5e7c449e4 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,19 +22,24 @@ 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)
@@ -52,12 +57,12 @@ 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})
@@ -126,3 +131,66 @@ class CacheDecoratorTest(testing.AppengineTestCase):
     a2 = A('http://test', 5)
     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, {})