Merge "Clean leave group hooks when unwatching."

2 files changed, 16 insertions, 0 deletions

diff --git a/tooz/coordination.py b/tooz/coordination.py
index 4307684..fdc80a0 100644
--- a/tooz/coordination.py
+++ b/tooz/coordination.py
@@ -322,6 +322,9 @@ class CoordinationDriver(object):
            group_id in self._group_members):
             del self._group_members[group_id]
 
+        if not self._hooks_leave_group[group_id]:
+            del self._hooks_leave_group[group_id]
+
     @abc.abstractmethod
     def watch_elected_as_leader(self, group_id, callback):
         """Call a function when member gets elected as leader.
diff --git a/tooz/tests/test_coordination.py b/tooz/tests/test_coordination.py
index f2181d4..2d93dbc 100644
--- a/tooz/tests/test_coordination.py
+++ b/tooz/tests/test_coordination.py
@@ -695,6 +695,19 @@ class TestAPI(tests.TestCaseSkipNotImplemented):
                           self._coord.unwatch_join_group,
                           self.group_id, lambda x: None)
 
+    def test_unwatch_leave_group(self):
+        # Create a group and add a leave_group callback
+        self._coord.create_group(self.group_id).get()
+        self.assertEqual(0, len(self._coord._hooks_leave_group))
+        self._coord.watch_leave_group(self.group_id, self._set_event)
+
+        # Ensure exactly one leave group hook exists
+        self.assertEqual(1, len(self._coord._hooks_leave_group[self.group_id]))
+
+        # Unwatch, and ensure no leave group hooks exist
+        self._coord.unwatch_leave_group(self.group_id, self._set_event)
+        self.assertEqual(0, len(self._coord._hooks_leave_group))
+
     def test_unwatch_leave_group_callback_not_found(self):
         self._coord.create_group(self.group_id).get()
         self.assertRaises(tooz.coordination.WatchCallbackNotFound,