From ff499604f0cfc5ce1b9e581ce3838ef943d78860 Mon Sep 17 00:00:00 2001
From: Oleksii Kliukin <oleksii.kliukin@zalando.de>
Date: Tue, 8 Sep 2015 16:04:54 +0200
Subject: [PATCH] Act on removal of initialization flag.

If initializer node suddenly dies before the initialization is complete,
other nodes should try to take over.

Fix some unittests for etcd and zookeeper and add couple of new ones.
---
 patroni/__init__.py     | 23 ++++++++++++-----------
 tests/test_etcd.py      |  4 ++++
 tests/test_zookeeper.py |  9 ++++++++-
 3 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/patroni/__init__.py b/patroni/__init__.py
index 46c030b6..b55c0dfc 100644
--- a/patroni/__init__.py
+++ b/patroni/__init__.py
@@ -55,17 +55,18 @@ class Patroni:
 
         # is data directory empty?
         if self.postgresql.data_directory_empty():
-            # racing to initialize
-            if self.ha.dcs.initialize():
-                try:
-                    self.postgresql.bootstrap()
-                except:
-                    # bail out and clean the initialize flag.
-                    self.cleanup_on_failed_initialization()
-                    raise
-                self.ha.dcs.take_leader()
-            else:
-                while True:
+            while True:
+                # racing to initialize
+                if self.ha.dcs.initialize():
+                    try:
+                        self.postgresql.bootstrap()
+                    except:
+                        # bail out and clean the initialize flag.
+                        self.cleanup_on_failed_initialization()
+                        raise
+                    self.ha.dcs.take_leader()
+                    break
+                else:
                     leader = self.ha.dcs.current_leader()
                     if leader and self.postgresql.bootstrap(leader):
                         break
diff --git a/tests/test_etcd.py b/tests/test_etcd.py
index b5c014f4..dbcd4b08 100644
--- a/tests/test_etcd.py
+++ b/tests/test_etcd.py
@@ -269,6 +269,10 @@ class TestEtcd(unittest.TestCase):
     def test_initialize(self):
         self.assertFalse(self.etcd.initialize())
 
+    def test_cancel_initializion(self):
+        self.etcd.client.delete = etcd_delete
+        self.assertFalse(self.etcd.cancel_initialization())
+
     def test_delete_leader(self):
         self.etcd.client.delete = etcd_delete
         self.assertFalse(self.etcd.delete_leader())
diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py
index d138e4b2..83c97b7f 100644
--- a/tests/test_zookeeper.py
+++ b/tests/test_zookeeper.py
@@ -71,6 +71,8 @@ class MockKazooClient:
             if self.leader:
                 return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
             return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+        elif path.endswith(patroni.zookeeper.ZooKeeper.initialize_key):
+            return 'foo'
 
     def get_children(self, path, watch=None, include_data=False):
         return ['foo', 'bar', 'buzz']
@@ -93,6 +95,8 @@ class MockKazooClient:
                 return
             self.leader = True
             raise Exception
+        elif path.endswith(patroni.zookeeper.ZooKeeper.initialize_key):
+            raise Exception
 
     def set_hosts(self, hosts, randomize_hosts=None):
         pass
@@ -146,9 +150,12 @@ class TestZooKeeper(unittest.TestCase):
         self.zk.touch_member('foo')
         self.zk.delete_leader()
 
-    def test_race(self):
+    def test_initialize(self):
         self.assertFalse(self.zk.initialize())
 
+    def test_cancel_initialization(self):
+        self.assertRaises(Exception, self.zk.cancel_initialization)
+
     def test_touch_member(self):
         self.zk.touch_member('new')
         self.zk.touch_member('exists')