Update behavior of lock to behave closer to redis lock

fakeredis.py

@@ -8,15 +8,16 @@
 from datetime import datetime, timedelta
 import operator
 import sys
-import threading
 import time
 import types
 import re
 import functools
+import threading
 from itertools import count, islice
+from uuid import uuid4
 
 import redis
-from redis.exceptions import ResponseError
+from redis.exceptions import ResponseError, LockError
 import redis.client
 
 try:
@@ -311,12 +312,15 @@ def _compile_pattern(pattern):
     return re.compile(regex, re.S)
 
 
+threading_lock = threading.Lock()
+
+
 class _Lock(object):
     def __init__(self, redis, name, timeout):
         self.redis = redis
         self.name = name
-        self.lock = threading.Lock()
-        redis.set(name, self, ex=timeout)
+        self.timeout = timeout
+        self.id = None
 
     def __enter__(self):
         self.acquire()
@@ -326,11 +330,27 @@ def __exit__(self, exc_type, exc_value, traceback):
         self.release()
 
     def acquire(self, blocking=True, blocking_timeout=None):
-        return self.lock.acquire(blocking)
+        token = str(uuid4())
+        with threading_lock:
+            acquired = bool(self.redis.set(self.name, token, nx=True, ex=self.timeout))
+        if not acquired and blocking:
+            raise ValueError('fakeredis can\'t do blocking locks')
+
+        if acquired:
+            self.id = token
+
+        return acquired
 
     def release(self):
-        self.lock.release()
-        self.redis.delete(self.name)
+        if self.id is None:
+            raise LockError("Cannot release an unlocked lock")
+
+        with threading_lock:
+            if _decode(self.redis.get(self.name)) != self.id:
+                raise LockError("Cannot release a lock that's no longer owned")
+
+            self.redis.delete(self.name)
+            self.id = None
 
 
 def _check_conn(func):

test_fakeredis.py

@@ -3916,6 +3916,45 @@ def test_lock(self):
             self.assertTrue(self.redis.exists('bar'))
         self.assertFalse(self.redis.exists('bar'))
 
+    def test_acquiring_lock_twice(self):
+        lock = self.redis.lock('foo')
+        self.assertTrue(lock.acquire(blocking=False))
+        self.assertFalse(lock.acquire(blocking=False))
+
+    def test_acquiring_lock_different_lock(self):
+        lock1 = self.redis.lock('foo')
+        lock2 = self.redis.lock('foo')
+        self.assertTrue(lock1.acquire(blocking=False))
+        self.assertFalse(lock2.acquire(blocking=False))
+
+    def test_acquiring_lock_different_lock_release(self):
+        lock1 = self.redis.lock('foo')
+        lock2 = self.redis.lock('foo')
+        self.assertTrue(lock1.acquire(blocking=False))
+        self.assertFalse(lock2.acquire(blocking=False))
+
+        # Test only releasing lock1 actually releases the lock
+        with self.assertRaises(redis.exceptions.LockError):
+            lock2.release()
+        self.assertFalse(lock2.acquire(blocking=False))
+        lock1.release()
+
+        # Locking with lock2 now has the lock
+        self.assertTrue(lock2.acquire(blocking=False))
+        self.assertFalse(lock1.acquire(blocking=False))
+
+    def test_nested_lock(self):
+        with self.redis.lock('bar'):
+            acquired = self.redis.lock('bar').acquire(blocking=False)
+            self.assertFalse(acquired)
+
+    def test_lock_no_longer_owned(self):
+        lock = self.redis.lock('bar')
+        lock.acquire()
+        self.redis.delete('bar')
+        with self.assertRaises(redis.exceptions.LockError):
+            lock.release()
+
 
 class DecodeMixin(object):
     decode_responses = True