Add transaction method to client

Author: tomlinford

mockredis/client.py

@@ -11,7 +11,7 @@
 
 from mockredis.clock import SystemClock
 from mockredis.lock import MockRedisLock
-from mockredis.exceptions import RedisError, ResponseError
+from mockredis.exceptions import RedisError, ResponseError, WatchError
 from mockredis.pipeline import MockRedisPipeline
 from mockredis.script import Script
 from mockredis.sortedset import SortedSet
@@ -75,6 +75,30 @@ def pipeline(self, transaction=True, shard_hint=None):
         """Emulate a redis-python pipeline."""
         return MockRedisPipeline(self, transaction, shard_hint)
 
+    def transaction(self, func, *watches, **kwargs):
+        """
+        Convenience method for executing the callable `func` as a transaction
+        while watching all keys specified in `watches`. The 'func' callable
+        should expect a single argument which is a Pipeline object.
+
+        Copied directly from redis-py.
+        """
+        shard_hint = kwargs.pop('shard_hint', None)
+        value_from_callable = kwargs.pop('value_from_callable', False)
+        watch_delay = kwargs.pop('watch_delay', None)
+        with self.pipeline(True, shard_hint) as pipe:
+            while 1:
+                try:
+                    if watches:
+                        pipe.watch(*watches)
+                    func_value = func(pipe)
+                    exec_value = pipe.execute()
+                    return func_value if value_from_callable else exec_value
+                except WatchError:
+                    if watch_delay is not None and watch_delay > 0:
+                        time.sleep(watch_delay)
+                    continue
+
     def watch(self, *argv, **kwargs):
         """
         Mock does not support command buffering so watch

mockredis/tests/test_pipeline.py

@@ -33,6 +33,30 @@ def test_pipeline_args(self):
         with self.redis.pipeline(transaction=False, shard_hint=None):
             pass
 
+    def test_transaction(self):
+        self.redis["a"] = 1
+        self.redis["b"] = 2
+        has_run = []
+
+        def my_transaction(pipe):
+            a_value = pipe.get("a")
+            assert a_value in (b"1", b"2")
+            b_value = pipe.get("b")
+            assert b_value == b"2"
+
+            # silly run-once code... incr's "a" so WatchError should be raised
+            # forcing this all to run again. this should incr "a" once to "2"
+            if not has_run:
+                self.redis.incr("a")
+                has_run.append(True)
+
+            pipe.multi()
+            pipe.set("c", int(a_value) + int(b_value))
+
+        result = self.redis.transaction(my_transaction, "a", "b")
+        eq_([True], result)
+        eq_(b"4", self.redis["c"])
+
     def test_set_and_get(self):
         """
         Pipeline execution returns the pipeline, not the intermediate value.