Implement Consul support

.travis.yml

@@ -9,11 +9,12 @@ addons:
   postgresql: "9.5"
 env:
   global:
-  - ETCDVERSION=2.3.1 ZKVERSION=3.4.6
+  - ETCDVERSION=2.3.2 ZKVERSION=3.4.6 CONSULVERSION=0.6.4
   matrix:
   - TEST_SUITE="python setup.py"
   - DCS="etcd" TEST_SUITE="behave"
   - DCS="exhibitor" TEST_SUITE="behave"
+  - DCS="consul" TEST_SUITE="behave"
 cache:
   directories:
   - $HOME/virtualenv/python2.7.9
@@ -24,6 +25,12 @@ install:
     set -e
 
     if [[ $TEST_SUITE == "behave" ]]; then
+        if [[ $DCS == "consul" ]]; then
+            curl -L https://releases.hashicorp.com/consul/${CONSULVERSION}/consul_${CONSULVERSION}_linux_amd64.zip \
+                | gunzip > consul
+            chmod +x consul
+        fi
+
         if [[ $DCS == "etcd" ]]; then
             curl -L https://github.com/coreos/etcd/releases/download/v${ETCDVERSION}/etcd-v${ETCDVERSION}-linux-amd64.tar.gz \
                 | tar xz -C . --strip=1 --wildcards --no-anchored etcd

README.rst

@@ -1,6 +1,6 @@
 |Build Status| |Coverage Status|
 
-Patroni: A Template for PostgreSQL HA with ZooKeeper or etcd
+Patroni: A Template for PostgreSQL HA with ZooKeeper, etcd or Consul
 ------------------------------------------------------------
 
 Patroni was previously known as Governor.
@@ -8,7 +8,7 @@ Patroni was previously known as Governor.
 *There are many ways to run high availability with PostgreSQL. Here, we
 present a template for you to create your own customized, high-availability 
 solution using Python and — for maximum accessibility — a distributed 
-configuration store like ZooKeeper or etcd.*
+configuration store like ZooKeeper, etcd or Consul.*
 
 Getting Started
 ---------------
@@ -61,12 +61,17 @@ For an example file, see ``postgres0.yml``. Regarding settings:
     -  *keyfile*: (optional) Specifies a file with the secret key in the PEM format.
 
 -  *etcd*:
-    -  *scope*: the relative path used on etcd's HTTP API for this deployment; makes it possible to run multiple HA deployments from a single etcd.
+    -  *scope*: the relative path used on etcd's HTTP API for this deployment; makes it possible to run multiple HA deployments from a single etcd cluster.
     -  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
     -  *host*: the host:port for the etcd endpoint.
 
+-  *consul*:
+    -  *scope*: the relative path used on Consul's HTTP API for this deployment; makes it possible to run multiple HA deployments from a single Consul cluster.
+    -  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
+    -  *host*: the host:port for the Consul endpoint.
+
 -  *zookeeper*:
-    -  *scope*: the relative path used on etcd's HTTP API for this deployment; makes it possible to run multiple HA deployments from a single etcd.
+    -  *scope*: the relative path used on ZooKeeper for this deployment; makes it possible to run multiple HA deployments from a single ZooKeeper cluster.
     -  *session\_timeout*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
     -  *reconnect\_timeout*: how long we should try to reconnect to ZooKeeper after a connection loss. After this timeout, assume that you no longer have a lock and restart in read-only mode.
     -  *hosts*: list of ZooKeeper cluster members in format: ['host1:port1', 'host2:port2', 'etc...']

features/environment.py

@@ -1,4 +1,5 @@
 import abc
+import consul
 import etcd
 import kazoo.client
 import kazoo.exceptions
@@ -128,13 +129,18 @@ class PatroniController(AbstractController):
             config['tags'] = tags
 
         if dcs != 'etcd':
-            etcd = config.pop('etcd')
-            config['zookeeper'] = {'scope': etcd['scope'], 'session_timeout': etcd['ttl'],
-                                   'reconnect_timeout': config['loop_wait']}
-            if dcs == 'exhibitor':
-                config['zookeeper']['exhibitor'] = {'hosts': ['127.0.0.1'], 'port': 8181}
+            dcs_config = config.pop('etcd')
+            dcs_config.pop('host')
+
+            if dcs == 'consul':
+                config[dcs] = dcs_config
             else:
-                config['zookeeper']['hosts'] = ['127.0.0.1:2181']
+                dcs_config.update({'session_timeout': dcs_config.pop('ttl'), 'reconnect_timeout': config['loop_wait']})
+                if dcs == 'exhibitor':
+                    dcs_config['exhibitor'] = {'hosts': ['127.0.0.1'], 'port': 8181}
+                else:
+                    dcs_config['hosts'] = ['127.0.0.1:2181']
+                config['zookeeper'] = dcs_config
 
         with open(patroni_config_path, 'w') as f:
             yaml.dump(config, f, default_flow_style=False)
@@ -196,6 +202,30 @@ class AbstractDcsController(AbstractController):
         """ clean all contents stored in the tree used for the tests """
 
 
+class ConsulController(AbstractDcsController):
+
+    def __init__(self, output_dir):
+        super(ConsulController, self).__init__('consul', tempfile.mkdtemp(), output_dir)
+        self._client = consul.Consul()
+
+    def _start(self):
+        return subprocess.Popen(['consul', 'agent', '-server', '-bootstrap', '-advertise=127.0.0.1',
+                                 '-data-dir', self._work_directory], stdout=self._log, stderr=subprocess.STDOUT)
+
+    def _is_running(self):
+        try:
+            return bool(self._client.status.leader())
+        except Exception:
+            return False
+
+    def query(self, key):
+        _, value = self._client.kv.get('{0}/{1}'.format(self._CLUSTER_NODE, key))
+        return value and value['Value'].decode('utf-8')
+
+    def cleanup_service_tree(self):
+        self._client.kv.delete(self._CLUSTER_NODE, recurse=True)
+
+
 class EtcdController(AbstractDcsController):
 
     """ handles all etcd related tasks, used for the tests setup and cleanup """
@@ -267,7 +297,8 @@ class ZooKeeperController(AbstractDcsController):
 
 class PatroniPoolController(object):
 
-    KNOWN_DCS = {'etcd': EtcdController, 'zookeeper': ZooKeeperController, 'exhibitor': ZooKeeperController}
+    KNOWN_DCS = {'consul': ConsulController, 'etcd': EtcdController,
+                 'zookeeper': ZooKeeperController, 'exhibitor': ZooKeeperController}
 
     def __init__(self):
         self._dcs = None

patroni/__init__.py

@@ -48,6 +48,9 @@ class Patroni(object):
         if 'zookeeper' in config:
             from patroni.zookeeper import ZooKeeper
             return ZooKeeper(name, config['zookeeper'])
+        if 'consul' in config:
+            from patroni.consul import Consul
+            return Consul(name, config['consul'])
         raise PatroniException('Can not find suitable configuration of distributed configuration store')
 
     def schedule_next_run(self):

patroni/consul.py

@@ -0,0 +1,244 @@
+from __future__ import absolute_import
+import logging
+import os
+import time
+import six
+
+from consul import ConsulException, NotFound, base, std
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
+from patroni.utils import sleep
+from requests.exceptions import RequestException
+
+logger = logging.getLogger(__name__)
+
+
+class ConsulError(DCSError):
+    pass
+
+
+class HTTPClient(std.HTTPClient):
+
+    def __init__(self, *args, **kwargs):
+        super(HTTPClient, self).__init__(*args, **kwargs)
+        self._patch_default_timeout()
+
+    def _patch_default_timeout(self):
+        # Set a default timeout for the `request.session.request` method, that is used
+        # internally  by the methods request.session.get, request.session.post and
+        # others. We monkey-patch here to avoid reimplementing each individual method from
+        # `std.HTTPClient`. By default, the timeout is not set. It means that a new
+        # session may hang almost indefinitely waiting for the server to respond,
+        # which is not what we want in Patroni.
+
+        request_func = getattr(self.session.request, '__func__' if six.PY3 else 'im_func')
+        defaults_attr_name = '__defaults__' if six.PY3 else 'func_defaults'
+        defaults = list(getattr(request_func, defaults_attr_name))
+        code = request_func.__code__ if six.PY3 else request_func.func_code
+        defaults[code.co_varnames[code.co_argcount - len(defaults):code.co_argcount].index('timeout')] = 5
+        setattr(request_func, defaults_attr_name, tuple(defaults))  # monkeypatching
+
+    def get(self, callback, path, params=None):
+        # The get function is overridden to handle a special case of it being called
+        # with an index and wait parameters. That form indicates that a user needs to
+        # wait for the given key to change its value, with a wait timeout supplied. We
+        # don't want our monkey-patched timeout to be less than the value of the wait
+        # parameter, therefore, we set it to either the value of wait or a default of 5 minutes.
+
+        if isinstance(params, dict) and 'index' in params:
+            timeout = (float(params['wait'][:-1]) if 'wait' in params else 300) + 1
+        else:
+            timeout = None
+        return callback(self.response(self.session.get(self.uri(path, params), verify=self.verify, timeout=timeout)))
+
+
+class ConsulClient(base.Consul):
+
+    @staticmethod
+    def connect(host, port, scheme, verify=True):
+        return HTTPClient(host, port, scheme, verify)
+
+
+def catch_consul_errors(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except (ConsulException, RequestException):
+            return False
+    return wrapper
+
+
+class Consul(AbstractDCS):
+
+    def __init__(self, name, config):
+        super(Consul, self).__init__(name, config)
+        self.ttl = int((config.get('ttl') or 30)/2)  # My experiments have shown that session expires after 2*ttl time
+        host, port = config.get('host', '127.0.0.1:8500').split(':')
+        self._client = ConsulClient(host=host, port=port)
+        self._scope = config['scope']
+        self._session = None
+        self._my_member_data = None
+        self.create_or_restore_session()
+
+    def create_or_restore_session(self):
+        while not self._session:
+            try:
+                _, member = self._client.kv.get(self.member_path)
+                self._session = (member or {}).get('Session')
+                if self.refresh_session():
+                    self._client.kv.delete(self.member_path)
+            except (ConsulException, RequestException):
+                logger.info('waiting on consul')
+                sleep(5)
+
+    def refresh_session(self):
+        """:returns: `!True` if it had to create new session"""
+        if self._session:
+            try:
+                return self._client.session.renew(self._session) is None
+            except NotFound:
+                self._session = None
+        if not self._session:
+            name = self._scope + '-' + self._name
+            try:
+                self._session = self._client.session.create(name=name, lock_delay=0, behavior='delete', ttl=self.ttl)
+            except (ConsulException, RequestException):
+                logger.exception('session.create')
+        if not self._session:
+            raise ConsulError('Failed to renew/create session')
+        return True
+
+    def client_path(self, path):
+        return super(Consul, self).client_path(path)[1:]
+
+    @staticmethod
+    def member(node):
+        return Member.from_node(node['ModifyIndex'], os.path.basename(node['Key']), node.get('Session'), node['Value'])
+
+    def _load_cluster(self):
+        try:
+            path = self.client_path('/')
+            _, results = self._client.kv.get(path, recurse=True)
+
+            if results is None:
+                raise NotFound
+
+            nodes = {}
+            for node in results:
+                node['Value'] = (node['Value'] or b'').decode('utf-8')
+                nodes[os.path.relpath(node['Key'], path)] = node
+
+            # get initialize flag
+            initialize = nodes.get(self._INITIALIZE)
+            initialize = initialize and initialize['Value']
+
+            # get last leader operation
+            last_leader_operation = nodes.get(self._LEADER_OPTIME)
+            last_leader_operation = 0 if last_leader_operation is None else int(last_leader_operation['Value'])
+
+            # get list of members
+            members = [self.member(n) for k, n in nodes.items() if k.startswith(self._MEMBERS) and k.count('/') == 1]
+
+            # get leader
+            leader = nodes.get(self._LEADER)
+            if leader and leader['Value'] == self._name and self._session != leader.get('Session', 'x'):
+                logger.info('I am leader but not owner of the session. Removing leader node')
+                self._client.kv.delete(self.leader_path, cas=leader['ModifyIndex'])
+                leader = None
+
+            if leader:
+                member = Member(-1, leader['Value'], None, {})
+                member = ([m for m in members if m.name == leader['Value']] or [member])[0]
+                leader = Leader(leader['ModifyIndex'], leader.get('Session'), member)
+
+            # failover key
+            failover = nodes.get(self._FAILOVER)
+            if failover:
+                failover = Failover.from_node(failover['ModifyIndex'], failover['Value'])
+
+            self._cluster = Cluster(initialize, leader, last_leader_operation, members, failover)
+        except NotFound:
+            self._cluster = Cluster(False, None, None, [], None)
+        except:
+            logger.exception('get_cluster')
+            raise ConsulError('Consul is not responding properly')
+
+    def touch_member(self, data, **kwargs):
+        create_member = self.refresh_session()
+        cluster = self.cluster
+        member = cluster and ([m for m in cluster.members if m.name == self._name] or [None])[0]
+        if create_member and member:
+            try:
+                self._client.kv.delete(self.member_path)
+            except Exception:
+                return False
+
+        if not create_member and member and data == self._my_member_data:
+            return True
+
+        try:
+            self._client.kv.put(self.member_path, data, acquire=self._session)
+            self._my_member_data = data
+            return True
+        except Exception:
+            logger.exception('touch_member')
+        return False
+
+    @catch_consul_errors
+    def attempt_to_acquire_leader(self):
+        ret = self._client.kv.put(self.leader_path, self._name, acquire=self._session)
+        if not ret:
+            logger.info('Could not take out TTL lock')
+        return ret
+
+    def take_leader(self):
+        return self.attempt_to_acquire_leader()
+
+    @catch_consul_errors
+    def set_failover_value(self, value, index=None):
+        return self._client.kv.put(self.failover_path, value, cas=index)
+
+    @catch_consul_errors
+    def write_leader_optime(self, last_operation):
+        return self._client.kv.put(self.leader_optime_path, last_operation)
+
+    @staticmethod
+    def update_leader():
+        return True
+
+    @catch_consul_errors
+    def initialize(self, create_new=True, sysid=''):
+        kwargs = {'cas': 0} if create_new else {}
+        return self._client.kv.put(self.initialize_path, sysid, **kwargs)
+
+    @catch_consul_errors
+    def cancel_initialization(self):
+        return self._client.kv.delete(self.initialize_path)
+
+    @catch_consul_errors
+    def delete_cluster(self):
+        return self._client.kv.delete(self.client_path(''), recurse=True)
+
+    @catch_consul_errors
+    def delete_leader(self):
+        cluster = self.cluster
+        if cluster and isinstance(cluster.leader, Leader) and cluster.leader.name == self._name:
+            return self._client.kv.delete(self.leader_path, cas=cluster.leader.index)
+
+    def watch(self, timeout):
+        cluster = self.cluster
+        if cluster and cluster.leader and cluster.leader.name != self._name and cluster.leader.index:
+            end_time = time.time() + timeout
+            while timeout >= 1:
+                try:
+                    idx, _ = self._client.kv.get(self.leader_path, index=cluster.leader.index, wait=str(timeout) + 's')
+                    return str(idx) != str(cluster.leader.index)
+                except (ConsulException, RequestException):
+                    logging.exception('watch')
+
+                timeout = end_time - time.time()
+
+        try:
+            return super(Consul, self).watch(timeout)
+        finally:
+            self.event.clear()

patroni/ctl.py

@@ -295,7 +295,7 @@ def query(
     if p_file is not None:
         command = p_file.read()
 
-    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
 
     cursor = None
     for _ in watching(w, watch, clear=False):
@@ -344,7 +344,7 @@ def query_member(cluster, cursor, member, role, command, connect_parameters=None
 @option_format
 @option_dcs
 def remove(config_file, cluster_name, fmt, dcs):
-    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
 
     output_members(cluster, cluster_name, fmt)
 
@@ -425,7 +425,7 @@ def ctl_load_config(cluster_name, config_file, dcs):
 @option_force
 @option_dcs
 def restart(cluster_name, member_names, config_file, dcs, force, role, p_any):
-    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
 
     role_names = [m.name for m in get_all_members(cluster, role)]
 
@@ -449,7 +449,7 @@ def restart(cluster_name, member_names, config_file, dcs, force, role, p_any):
 @option_force
 @option_dcs
 def reinit(cluster_name, member_names, config_file, dcs, force):
-    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
     empty_post_to_members(cluster, member_names, force, 'reinitialize')
 
 
@@ -470,7 +470,7 @@ def failover(config_file, cluster_name, master, candidate, force, dcs, scheduled
         If so, we trigger a failover and keep the client up to date.
     """
 
-    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
 
     if cluster.leader is None:
         raise PatroniCtlException('This cluster has no master')
@@ -541,7 +541,7 @@ def failover(config_file, cluster_name, master, candidate, force, dcs, scheduled
         else:
             click.echo('Failover failed, details: {0}, {1}'.format(r.status_code, r.text))
             return
-    except:
+    except Exception:
         logging.exception(r)
         logging.warning('Failing over to DCS')
         click.echo(timestamp() + ' Could not failover using Patroni api, falling back to DCS')

postgres0.yml

@@ -12,6 +12,10 @@ etcd:
   ttl: *ttl
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#consul:
+#  scope: *scope
+#  ttl: *ttl
+#  host: 127.0.0.1:8500
 #zookeeper:
 #  scope: *scope
 #  session_timeout: *ttl

postgres1.yml

@@ -12,6 +12,10 @@ etcd:
   ttl: *ttl
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#consul:
+#  scope: *scope
+#  ttl: *ttl
+#  host: 127.0.0.1:8500
 #zookeeper:
 #  scope: *scope
 #  session_timeout: *ttl

postgres2.yml

@@ -12,6 +12,10 @@ etcd:
   ttl: *ttl
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#consul:
+#  scope: *scope
+#  ttl: *ttl
+#  host: 127.0.0.1:8500
 #zookeeper:
 #  scope: *scope
 #  session_timeout: *ttl

requirements.txt

@@ -5,6 +5,7 @@ requests
 six >= 1.7
 kazoo>=2.2.1
 python-etcd==0.4.3
+python-consul==0.6.0
 click>=4.1
 prettytable>=0.7
 tzlocal

tests/test_consul.py

@@ -0,0 +1,128 @@
+import consul
+import unittest
+
+from patroni.dcs import AbstractDCS
+from mock import Mock, patch
+from patroni.consul import Cluster, Consul, ConsulError, ConsulException, HTTPClient, NotFound
+from test_etcd import SleepException
+
+
+def kv_get(self, key, **kwargs):
+    if key == 'service/test/members/postgresql1':
+        return '1', {'Session': 'fd4f44fe-2cac-bba5-a60b-304b51ff39b7'}
+    if key == 'service/test/':
+        return None, None
+    if key == 'service/good/leader':
+        return '1', None
+    if key == 'service/good/':
+        return ('6429',
+                [{'CreateIndex': 1334, 'Flags': 0, 'Key': key + 'failover', 'LockIndex': 0,
+                  'ModifyIndex': 1334, 'Value': b''},
+                 {'CreateIndex': 1334, 'Flags': 0, 'Key': key + 'initialize', 'LockIndex': 0,
+                  'ModifyIndex': 1334, 'Value': b'postgresql0'},
+                 {'CreateIndex': 2621, 'Flags': 0, 'Key': key + 'leader', 'LockIndex': 1,
+                  'ModifyIndex': 2621, 'Session': 'fd4f44fe-2cac-bba5-a60b-304b51ff39b7', 'Value': b'postgresql1'},
+                 {'CreateIndex': 6156, 'Flags': 0, 'Key': key + 'members/postgresql0', 'LockIndex': 1,
+                  'ModifyIndex': 6156, 'Session': '782e6da4-ed02-3aef-7963-99a90ed94b53',
+                  'Value': ('postgres://replicator:[email protected]:5432/postgres' +
+                            '?application_name=http://127.0.0.1:8008/patroni').encode('utf-8')},
+                 {'CreateIndex': 2630, 'Flags': 0, 'Key': key + 'members/postgresql1', 'LockIndex': 1,
+                  'ModifyIndex': 2630, 'Session': 'fd4f44fe-2cac-bba5-a60b-304b51ff39b7',
+                  'Value': ('postgres://replicator:[email protected]:5433/postgres' +
+                            '?application_name=http://127.0.0.1:8009/patroni').encode('utf-8')},
+                 {'CreateIndex': 1085, 'Flags': 0, 'Key': key + 'optime/leader', 'LockIndex': 0,
+                  'ModifyIndex': 6429, 'Value': b'4496294792'}])
+    raise ConsulException
+
+
+class TestHTTPClient(unittest.TestCase):
+
+    def test_get(self):
+        self.client = HTTPClient('127.0.0.1', '8500', 'http', False)
+        self.client.session.get = Mock()
+        self.client.get(Mock(), '')
+        self.client.get(Mock(), '', {'wait': '1s', 'index': 1})
+
+
+@patch.object(consul.Consul.KV, 'get', kv_get)
+class TestConsul(unittest.TestCase):
+
+    @patch.object(consul.Consul.Session, 'create', Mock(return_value='fd4f44fe-2cac-bba5-a60b-304b51ff39b7'))
+    @patch.object(consul.Consul.Session, 'renew', Mock(side_effect=NotFound))
+    @patch.object(consul.Consul.KV, 'get', kv_get)
+    @patch.object(consul.Consul.KV, 'delete', Mock())
+    def setUp(self):
+        self.c = Consul('postgresql1', {'ttl': 30, 'scope': 'test', 'host': 'localhost:1'})
+        self.c._base_path = '/service/good'
+        self.c._load_cluster()
+
+    @patch('time.sleep', Mock(side_effect=SleepException))
+    def test_create_or_restore_session(self):
+        self.c._session = None
+        self.assertRaises(SleepException, self.c.create_or_restore_session)
+
+    @patch.object(consul.Consul.Session, 'renew', Mock(side_effect=NotFound))
+    @patch.object(consul.Consul.Session, 'create', Mock(side_effect=ConsulException))
+    def test_referesh_session(self):
+        self.c._session = '1'
+        self.c._name = ''
+        self.assertRaises(ConsulError, self.c.refresh_session)
+
+    @patch.object(consul.Consul.KV, 'delete', Mock())
+    def test_get_cluster(self):
+        self.c._base_path = '/service/test'
+        self.assertIsInstance(self.c.get_cluster(), Cluster)
+        self.assertIsInstance(self.c.get_cluster(), Cluster)
+        self.c._base_path = '/service/fail'
+        self.assertRaises(ConsulError, self.c.get_cluster)
+        self.c._base_path = '/service/good'
+        self.c._session = 'fd4f44fe-2cac-bba5-a60b-304b51ff39b8'
+        self.assertIsInstance(self.c.get_cluster(), Cluster)
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(side_effect=[ConsulException, True, True]))
+    @patch.object(consul.Consul.KV, 'put', Mock(side_effect=[True, ConsulException]))
+    def test_touch_member(self):
+        self.c.refresh_session = Mock(return_value=True)
+        self.c.touch_member('balbla')
+        self.c.touch_member('balbla')
+        self.c.touch_member('balbla')
+        self.c.refresh_session = Mock(return_value=False)
+        self.c.touch_member('balbla')
+
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=False))
+    def test_take_leader(self):
+        self.c.take_leader()
+
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))
+    def test_set_failover_value(self):
+        self.c.set_failover_value('')
+
+    @patch.object(consul.Consul.KV, 'put', Mock(side_effect=ConsulException))
+    def test_write_leader_optime(self):
+        self.c.write_leader_optime('')
+
+    def test_update_leader(self):
+        self.c.update_leader()
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    def test_delete_leader(self):
+        self.c.delete_leader()
+
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))
+    def test_initialize(self):
+        self.c.initialize()
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    def test_cancel_initialization(self):
+        self.c.cancel_initialization()
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    def test_delete_cluster(self):
+        self.c.delete_cluster()
+
+    @patch.object(AbstractDCS, 'watch', Mock())
+    def test_watch(self):
+        self.c._name = ''
+        self.c.watch(1)
+        with patch.object(consul.Consul.KV, 'get', Mock(side_effect=ConsulException)):
+            self.c.watch(1)

tests/test_patroni.py

@@ -8,6 +8,7 @@ import yaml
 from mock import Mock, patch
 from patroni.api import RestApiServer
 from patroni.async_executor import AsyncExecutor
+from patroni.consul import Consul
 from patroni.etcd import Etcd
 from patroni import Patroni, PatroniException, main as _main
 from patroni.zookeeper import ZooKeeper
@@ -41,8 +42,10 @@ class TestPatroni(unittest.TestCase):
                 self.p = Patroni(config)
 
     @patch('patroni.zookeeper.KazooClient', MockKazooClient())
+    @patch.object(Consul, 'create_or_restore_session', Mock())
     def test_get_dcs(self):
         self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
+        self.assertIsInstance(self.p.get_dcs('', {'consul': {'scope': '', 'hosts': '127.0.0.1:1'}}), Consul)
         self.assertRaises(PatroniException, self.p.get_dcs, '', {})
 
     @patch('time.sleep', Mock(side_effect=SleepException))