From 924ca01f14bbaba2c5b9f15802760ef8636e033e Mon Sep 17 00:00:00 2001
From: Taylor Hakes <taylor@cedar.com>
Date: Mon, 17 Sep 2018 16:38:19 -0400
Subject: [PATCH] Update behavior of lock to behave closer to redis lock

---
 fakeredis.py      | 19 +++++++++++++------
 test_fakeredis.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/fakeredis.py b/fakeredis.py
index 921e3ff..c4dd7bf 100644
--- a/fakeredis.py
+++ b/fakeredis.py
@@ -8,12 +8,12 @@
 from datetime import datetime, timedelta
 import operator
 import sys
-import threading
 import time
 import types
 import re
 import functools
 from itertools import count, islice
+from uuid import uuid4
 
 import redis
 from redis.exceptions import ResponseError
@@ -315,8 +315,8 @@ 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 = str(uuid4())
 
     def __enter__(self):
         self.acquire()
@@ -326,11 +326,18 @@ def __exit__(self, exc_type, exc_value, traceback):
         self.release()
 
     def acquire(self, blocking=True, blocking_timeout=None):
-        return self.lock.acquire(blocking)
+        owner_id = _decode(self.redis.get(self.name))
+        can_acquire = not owner_id or owner_id == self.id
+        if can_acquire:
+            self.redis.set(self.name, self.id, ex=self.timeout)
+        elif blocking:
+            raise ValueError('fakeredis can\'t do blocking locks')
+
+        return can_acquire
 
     def release(self):
-        self.lock.release()
-        self.redis.delete(self.name)
+        if _decode(self.redis.get(self.name)) == self.id:
+            self.redis.delete(self.name)
 
 
 def _check_conn(func):
diff --git a/test_fakeredis.py b/test_fakeredis.py
index 715559c..14400ce 100644
--- a/test_fakeredis.py
+++ b/test_fakeredis.py
@@ -3916,6 +3916,37 @@ 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.assertTrue(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
+        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)
+
 
 class DecodeMixin(object):
     decode_responses = True