Accept blocking argument in lock's context manager

This patch adds the posibility of passing the blocking argument to the
context manager to be used when acquiring the lock.

Previous behavior will be preserved, so now we will be able to use the
context manager in any of the following ways:

    with lock(10.0):
    with lock(False):
    with lock():
    with lock:

The last 2 are equivalent.

Change-Id: I6d2f486ac45078ec4ce77de52b7997f2122d4dc8

tooz/locking.py

@@ -22,6 +22,18 @@ import weakref
 from tooz import coordination
 
 
+class _LockProxy(object):
+    def __init__(self, lock, blocking=True):
+        self.lock = lock
+        self.blocking = blocking
+
+    def __enter__(self):
+        return self.lock.__enter__(self.blocking)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.lock.__exit__(exc_type, exc_val, exc_tb)
+
+
 @six.add_metaclass(abc.ABCMeta)
 class Lock(object):
     def __init__(self, name):
@@ -33,8 +45,11 @@ class Lock(object):
     def name(self):
         return self._name
 
-    def __enter__(self):
-        acquired = self.acquire()
+    def __call__(self, blocking=True):
+        return _LockProxy(self, blocking)
+
+    def __enter__(self, blocking=True):
+        acquired = self.acquire(blocking)
         if not acquired:
             msg = u'Acquiring lock %s failed' % self.name
             raise coordination.LockAcquireFailed(msg)

tooz/tests/test_coordination.py

@@ -699,6 +699,47 @@ class TestAPI(testscenarios.TestWithScenarios,
         with lock as returned_lock:
             self.assertEqual(lock, returned_lock)
 
+    def test_lock_context_manager_acquire_no_argument(self):
+        name = self._get_random_uuid()
+        lock1 = self._coord.get_lock(name)
+        lock2 = self._coord.get_lock(name)
+        with lock1():
+            self.assertFalse(lock2.acquire(blocking=False))
+
+    def test_lock_context_manager_acquire_argument_return_value(self):
+        name = self._get_random_uuid()
+        blocking_value = 10.12
+        lock = self._coord.get_lock(name)
+        with lock(blocking_value) as returned_lock:
+            self.assertEqual(lock, returned_lock)
+
+    def test_lock_context_manager_acquire_argument_release_within(self):
+        name = self._get_random_uuid()
+        blocking_value = 10.12
+        lock = self._coord.get_lock(name)
+        with lock(blocking_value) as returned_lock:
+            self.assertTrue(returned_lock.release())
+
+    def test_lock_context_manager_acquire_argument(self):
+        name = self._get_random_uuid()
+        blocking_value = 10.12
+        lock = self._coord.get_lock(name)
+        with mock.patch.object(lock, 'acquire', wraps=True, autospec=True) as \
+                mock_acquire:
+            with lock(blocking_value):
+                mock_acquire.assert_called_once_with(blocking_value)
+
+    def test_lock_context_manager_acquire_argument_timeout(self):
+        name = self._get_random_uuid()
+        lock1 = self._coord.get_lock(name)
+        lock2 = self._coord.get_lock(name)
+        with lock1:
+            try:
+                with lock2(False):
+                    self.fail('Lock acquire should have failed')
+            except tooz.coordination.LockAcquireFailed:
+                pass
+
     def test_get_lock_locked_twice(self):
         name = self._get_random_uuid()
         lock = self._coord.get_lock(name)