Add some tests for zoop.lock

Rakefile

@@ -1,5 +1,5 @@
 
-PROJ = "rpc"
+PROJ = "zoop"
 
 task :test, :python do |t, args|
   p "Running unit tests for #{PROJ}"

test/test_lock.py

@@ -6,7 +6,7 @@
 if sys.version_info < (2, 7):
     import unittest2 as unittest
 
-from mock import Mock
+from mock import patch, Mock
 import zookeeper
 
 from zoop import lock
@@ -27,13 +27,56 @@ def test_init(self):
         self.zk.create.assert_any_call('/lockz/foolock')
         self.zk.create.assert_any_call('/lockz')
 
+    def test_contextmanager(self):
+        "Can we use it as a contextmanager?"
+        with patch.object(self.lk, 'acquire') as Pac:
+            with patch.object(self.lk, 'release') as Prel:
+                with self.lk:
+                    Pac.assert_called_once_with()
+                Prel.assert_called_once_with()
+
+    def test_revoked(self):
+        """ Predicatise a list """
+        self.lk.tlocal.revoked.append(True)
+        self.assertEqual(True, self.lk.revoked)
+
+    def test_not_revoked(self):
+        """ We haven't been asked to revoke yet. """
+        self.assertEqual(False, self.lk.revoked)
+
     def test_acquire(self):
         "Acquire the lock"
-        # self.lk.acquire()
-        # self.zk.create.assert_called_once_with('/zooplocks/barlock/baselock-',
-        #                                        value="0",
-        #                                        flags=zookeeper.SEQUENCE)
-        # !!! Test this
+        # !!! This is brittle.
+        self.zk.get.return_value = 'got'
+        self.zk.create.return_value = '/lockz/foolock/lock-00000001'
+        self.zk.get_children.return_value = ['lock-00000001']
+        self.assertEqual(True, self.lk.acquire())
+
+    def test_create_waitnode(self):
+        "Create a wait node."
+        self.zk.create.return_value = '/zooplocks/barlock/baselock-00000001'
+        self.zk.get.return_value = 'got'
+
+        nodepath, keynode = self.lk._create_waitnode()
+        self.zk.create.assert_called_once_with('/zooplocks/barlock/baselock-',
+                                               value = '0',
+                                               flags = zookeeper.SEQUENCE)
+        self.assertEqual('/zooplocks/barlock/baselock-00000001', nodepath)
+        self.assertEqual('baselock-00000001', keynode)
+
+    def test_has_lock(self):
+        """ Do we have the lock """
+        cases = [
+            ((True, None), ('baselock-0001', ['baselock-0001', 'baselock-0002'])),
+            ((False, ['baselock-0001']), ('baselock-0002', ['baselock-0001', 'baselock-0002']))
+            ]
+        for expected, arg in cases:
+            actual = self.lk.has_lock(*arg)
+            self.assertEqual(expected, actual)
+
+    def test_release(self):
+        "Can we release the lock?"
+        self.assertEqual(True, self.lk.release())
 
 class LockTestCase(unittest.TestCase):
     def setUp(self):

zoop/lock.py

@@ -53,6 +53,7 @@ def __init__(self, handle, name, root='/zooplocks'):
         self.tlocal = threading.local()
         self.tlocal.revoked = []
         self.tlocal.locking = None
+        self.tlocal.acquired = False
         if not self.zk.exists(root):
             self.zk.create(root)
         if not self.zk.exists(self.path):
@@ -89,55 +90,53 @@ def acquire(self, timeout=None):
         Arguments:
         - `timeout`: int
 
-        Return: None
+        Return: bool - whether we acquired the Lock or not
         Exceptions: None
         """
         # This implementation is based upon the Mozilla Services
         # ztools lock at https://github.com/mozilla-services/zktools
+        # with some additional encapsulation and error handling added.
         self.tlocal.revoked = []
 
         nodepath, keyname = self._create_waitnode()
-
-        acquired = False
         cv = threading.Event()
 
         def lockwatch(handle, etype, state, path):
             cv.set()
 
         tstart = time.time()
-
-        while not acquired:
+        frist = True
+        while not self.tlocal.acquired:
             cv.clear()
 
-            if timeout is not None and time.time() - tstart > timeout:
+            if not frist and timeout is not None and time.time() - tstart > timeout:
                 try:
                     self.zk.delete(nodepath)
                 except exceptions.NoNodeError:
                     pass
                 return False
+            frist = False
 
             kids = self.zk.get_children(self.path)
-            # This sort isn't used here, but in has_lock()
-            kids.sort(key=lambda val: val[val.rfind('-') + 1:])
 
             if len(kids) == 0 or not keyname in kids:
                 # Only really for connection issues
                 nodepath, keyname = self._create_waitnode()
                 continue
 
-            acquired, blocking = self.has_lock(keyname, kids)
-            if acquired:
+            self.acquired, blocking = self.has_lock(keyname, kids)
+            if self.acquired:
                 break
 
             last_blocker = join(self.path, blocking[-1])
             if not self.zk.exists(last_blocker, lockwatch):
-                continue # Wait - what?
+                continue # Already free
 
             if timeout is not None:
                 cv.wait(timeout - (time.time() - tstart))
 
         self.tlocal.lock_node = nodepath
-        return
+        return True
 
     def _create_waitnode(self):
         """
@@ -196,6 +195,8 @@ def has_lock(self, keypath, locknodes):
                 or None.
         Exceptions: None
         """
+        locknodes.sort(key=lambda val: val[val.rfind('-') + 1:])
+
         if keypath == locknodes[0]:
             return True, None
         return False, locknodes[:locknodes.index(keypath)]
@@ -208,12 +209,13 @@ def release(self):
         Exceptions: None
         """
         self.tlocal.revoked = []
+        self.acquired = False
         try:
             self.zk.delete(self.tlocal.lock_node)
             del self.tlocal.lock_node
         except (zookeeper.NoNodeException, AttributeError):
             pass # We never had the Lock!
-        return
+        return True
 
 class Lock(BaseLock):
     """