From bfa9c4adfbe4f0089e749cc997c71d0bdce0a0f9 Mon Sep 17 00:00:00 2001
From: Manish Singh <yosh@yosh.org>
Date: Sat, 6 Aug 2016 09:16:29 -0700
Subject: [PATCH] Update lock methods to match current redis-py

---
 mockredis/client.py |  9 +++++++--
 mockredis/lock.py   | 14 ++++++++++++--
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/mockredis/client.py b/mockredis/client.py
index 05c2439..387836f 100644
--- a/mockredis/client.py
+++ b/mockredis/client.py
@@ -74,9 +74,14 @@ def ping(self):
 
     # Transactions Functions #
 
-    def lock(self, key, timeout=0, sleep=0):
+    def lock(self, name, timeout=None, sleep=0.1, blocking_timeout=None,
+             lock_class=None, thread_local=True):
         """Emulate lock."""
-        return MockRedisLock(self, key, timeout, sleep)
+        if lock_class is None:
+            lock_class = MockRedisLock
+        return lock_class(self, name, timeout=timeout, sleep=sleep,
+                          blocking_timeout=blocking_timeout,
+                          thread_local=thread_local)
 
     def pipeline(self, transaction=True, shard_hint=None):
         """Emulate a redis-python pipeline."""
diff --git a/mockredis/lock.py b/mockredis/lock.py
index ddef603..c817502 100644
--- a/mockredis/lock.py
+++ b/mockredis/lock.py
@@ -4,7 +4,8 @@ class MockRedisLock(object):
     to allow testing without a real redis server.
     """
 
-    def __init__(self, redis, name, timeout=None, sleep=0.1):
+    def __init__(self, redis, name, timeout=None, sleep=0.1,
+                 blocking=True, blocking_timeout=None, thread_local=True):
         """Initialize the object."""
 
         self.redis = redis
@@ -12,6 +13,9 @@ def __init__(self, redis, name, timeout=None, sleep=0.1):
         self.acquired_until = None
         self.timeout = timeout
         self.sleep = sleep
+        self.blocking = blocking
+        self.blocking_timeout = blocking_timeout
+        self.thread_local = bool(thread_local)
 
     def acquire(self, blocking=True):  # pylint: disable=R0201,W0613
         """Emulate acquire."""
@@ -23,8 +27,14 @@ def release(self):   # pylint: disable=R0201
 
         return
 
+    def extend(self, additional_time):  # pylint: disable=R0201,W0613
+        """Emulate extend."""
+
+        return
+
     def __enter__(self):
-        return self.acquire()
+        self.acquire(blocking=True)
+        return self
 
     def __exit__(self, exc_type, exc_value, traceback):
         self.release()