Watch for changes of leader key

tests/test_etcd.py

@@ -3,6 +3,7 @@ import dns.resolver
 import etcd
 import json
 import requests
+import urllib3
 import socket
 import time
 import unittest
@@ -61,6 +62,20 @@ def requests_get(url, **kwargs):
     return response
 
 
+def etcd_watch(key, index=None, timeout=None, recursive=None):
+    print ('watch', key, index, timeout)
+    if timeout == 1:
+        raise urllib3.exceptions.TimeoutError
+    elif timeout == 5:
+        return etcd.EtcdResult('delete', {})
+    elif timeout == 10:
+        raise etcd.EtcdException
+    elif index == 20729:
+        return etcd.EtcdResult('set', {'value': 'postgresql1', 'modifiedIndex': index + 1})
+    elif index == 20731:
+        return etcd.EtcdResult('set', {'value': 'postgresql2', 'modifiedIndex': index + 1})
+
+
 def etcd_write(key, value, **kwargs):
     if key == '/service/test/leader':
         if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True):
@@ -237,3 +252,12 @@ class TestEtcd(unittest.TestCase):
     def test_delete_leader(self):
         self.etcd.client.delete = etcd_delete
         self.assertFalse(self.etcd.delete_leader())
+
+    def test_sleep(self):
+        self.etcd.client.watch = etcd_watch
+        self.etcd.sleep(100)
+        self.etcd.get_cluster()
+        self.etcd.sleep(1)
+        self.etcd.sleep(5)
+        self.etcd.sleep(10)
+        self.etcd.sleep(100)

tests/test_patroni.py

@@ -14,7 +14,7 @@ from helpers.zookeeper import ZooKeeper
 from mock import Mock, patch
 from patroni import Patroni, main
 from six.moves import BaseHTTPServer
-from test_etcd import Client, etcd_read, etcd_write
+from test_etcd import Client, etcd_read, etcd_write, etcd_watch
 from test_ha import true, false
 from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
 from test_zookeeper import MockKazooClient
@@ -104,6 +104,7 @@ class TestPatroni(unittest.TestCase):
         self.p.touch_member = self.touch_member
         self.p.ha.state_handler.sync_replication_slots = time_sleep
         self.p.ha.dcs.client.read = etcd_read
+        self.p.ha.dcs.sleep = time_sleep
         self.assertRaises(SleepException, self.p.run)
         self.p.ha.state_handler.is_leader = lambda: False
         self.p.api.start = nop