Merge pull request #10 from zalando/feature/python-etcd

python-etcd

helpers/api.py

@@ -1,3 +1,4 @@
+import fcntl
 import json
 import logging
 import psycopg2
@@ -55,6 +56,7 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
         host, port = config['listen'].split(':')
         HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
         Thread.__init__(self, target=self.serve_forever)
+        self._set_fd_cloexec(self.socket)
         self.patroni = patroni
         self.daemon = True
 
@@ -64,3 +66,8 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
         ret = [r for r in cursor]
         cursor.close()
         return ret
+
+    @staticmethod
+    def _set_fd_cloexec(fd):
+        flags = fcntl.fcntl(fd, fcntl.F_GETFD)
+        fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)

helpers/etcd.py

@@ -1,4 +1,7 @@
+from __future__ import absolute_import
+import etcd
 import logging
+import os
 import random
 import requests
 import socket
@@ -16,49 +19,37 @@ class EtcdError(DCSError):
     pass
 
 
-class EtcdConnectionFailed(EtcdError):
-    pass
-
-
-class Client:
-
-    API_VERSION = 'v2'
+class Client(etcd.Client):
 
     def __init__(self, config):
+        super(Client, self).__init__(read_timeout=5)
         self._config = config
-        self.timeout = 5
-        self._base_uri = None
-        self._members_cache = []
-        self.load_members()
+        self._load_machines_cache()
+        self._allow_reconnect = True
 
-    def client_url(self, path):
-        return self._base_uri + path
-
-    def _next_server(self):
-        self._base_uri = None
+    @property
+    def machines(self):
+        """Original `machines` method(property) of `etcd.Client` class raise exception
+        when it failed to get list of etcd cluster members. This method is being called
+        only when request failed on one of the etcd members during `api_execute` call.
+        For us it's more important to execute original request rather then get new
+        topology of etcd cluster. So we will catch this exception and return valid list
+        of machines with setting flag `self._update_machines_cache` to `!True`.
+        Later, during next `api_execute` call we will forcefully update machines_cache"""
         try:
-            self._base_uri = self._members_cache.pop()
-        except IndexError:
-            logger.error('Members cache is empty, can not retry.')
-            raise EtcdConnectionFailed('No more members in the cluster')
-        else:
-            logger.info('Selected new etcd server %s', self._base_uri)
+            ret = super(Client, self).machines
+            random.shuffle(ret)
+            return ret
+        except etcd.EtcdException:
+            if self._update_machines_cache:  # We are updating machines_cache
+                raise  # This exception is fatal, we should re-raise it.
+            self._update_machines_cache = True
+            return [self._base_uri]
 
-    def _get(self, path):
-        response = None
-        while response is None:
-            uri = self.client_url(path)
-            try:
-                logger.info('GET %s', uri)
-                response = requests.get(uri, timeout=self.timeout)
-            except RequestException:
-                self._next_server()
-
-        logger.debug([response.status_code, response.content])
-        try:
-            return response.json(), response.status_code
-        except (TypeError, ValueError):
-            raise EtcdError('Bad response from %s: %s' % (uri, response.content))
+    def api_execute(self, path, method, **kwargs):
+        # Update machines_cache if previous attempt of update has failed
+        self._update_machines_cache and self._load_machines_cache()
+        return super(Client, self).api_execute(path, method, **kwargs)
 
     @staticmethod
     def get_srv_record(host):
@@ -68,120 +59,74 @@ class Client:
             logger.exception('Can not resolve SRV for %s', host)
         return []
 
-    @staticmethod
-    def get_peers_urls_from_dns(host):
-        return ['http://{}:{}'.format(h, p) for h, p in Client.get_srv_record(host)]
+    def _get_machines_cache_from_srv(self, discovery_srv):
+        """Fetch list of etcd-cluster member by resolving _etcd-server._tcp. SRV record.
+        This record should contain list of host and peer ports which could be used to run
+        'GET http://{host}:{port}/members' request (peer protocol)"""
 
-    @staticmethod
-    def get_client_urls_from_dns(addr):
-        host, port = addr.split(':')
         ret = []
+        for host, port in self.get_srv_record(discovery_srv):
+            url = '{}://{}:{}/members'.format(self._protocol, host, port)
+            try:
+                response = requests.get(url)
+                if response.ok:
+                    for member in response.json():
+                        ret.extend(member['clientURLs'])
+                    break
+            except RequestException:
+                logger.exception('GET %s', url)
+        return list(set(ret))
+
+    def _get_machines_cache_from_dns(self, addr):
+        """One host might be resolved into multiple ip addresses. We will make list out of it"""
+
+        ret = []
+        host, port = addr.split(':')
         try:
             for r in set(socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)):
-                ret.append('http://{}:{}/{}'.format(r[4][0], r[4][1], Client.API_VERSION))
+                ret.append('{}://{}:{}'.format(self._protocol, r[4][0], r[4][1]))
         except socket.error:
             logger.exception('Can not resolve %s', host)
-        return list(set(ret)) if ret else ['http://{}:{}/{}'.format(host, port, Client.API_VERSION)]
+        return list(set(ret)) if ret else ['{}://{}:{}'.format(self._protocol, host, port)]
 
-    def load_members(self):
-        load_from_srv = False
-        if not self._base_uri:
-            if 'discovery_srv' not in self._config and 'host' not in self._config:
-                raise Exception('Neither discovery_srv nor host are defined in etcd section of config')
+    def _load_machines_cache(self):
+        """This method should fill up `_machines_cache` from scratch.
+        It could happen only in two cases:
+        1. During class initialization
+        2. When all etcd members failed"""
 
-            if 'discovery_srv' in self._config:
-                load_from_srv = True
-                self._members_cache = self.get_peers_urls_from_dns(self._config['discovery_srv'])
+        self._update_machines_cache = True
 
-            if not self._members_cache and 'host' in self._config:
-                load_from_srv = False
-                self._members_cache = self.get_client_urls_from_dns(self._config['host'])
+        if 'discovery_srv' not in self._config and 'host' not in self._config:
+            raise Exception('Neither discovery_srv nor host are defined in etcd section of config')
 
-            self._next_server()
+        self._machines_cache = []
 
-        response, status_code = self._get('/members')
-        if status_code != 200:
-            self._base_uri = None
-            raise EtcdError('Got response with code=%s from %s' % (status_code, self._base_uri))
+        if 'discovery_srv' in self._config:
+            self._machines_cache = self._get_machines_cache_from_srv(self._config['discovery_srv'])
 
-        members_cache = []
+        if not self._machines_cache and 'host' in self._config:
+            self._machines_cache = self._get_machines_cache_from_dns(self._config['host'])
+
+        # Can not bootstrap list of etcd-cluster members, giving up
+        if not self._machines_cache:
+            raise etcd.EtcdException
+
+        # After filling up initial list of machines_cache we should ask etcd-cluster about actual list
+        self._base_uri = self._machines_cache.pop(0)
+        self._machines_cache = self.machines
+        self._base_uri in self._machines_cache and self._machines_cache.remove(self._base_uri)
+
+        self._update_machines_cache = False
+
+
+def catch_etcd_errors(func):
+    def wrapper(*args, **kwargs):
         try:
-            for member in response if load_from_srv else response['members']:
-                members_cache.extend([m + '/' + self.API_VERSION for m in member['clientURLs']])
-        except:
-            self._base_uri = None
-            raise EtcdError('Got invalid response from %s: %s' % (self._base_uri, response))
-
-        self._members_cache = list(set(members_cache))
-        random.shuffle(self._members_cache)
-        if load_from_srv:
-            self._next_server()
-        else:
-            try:
-                self._members_cache.remove(self._base_uri)
-            except ValueError:
-                pass
-
-    def get(self, path):
-        if not self._base_uri:
-            self.load_members()
-        old_base_uri = self._base_uri
-        try:
-            return self._get(path)
-        finally:
-            if self._base_uri != old_base_uri:
-                try:
-                    self.load_members()
-                except EtcdError:
-                    logger.exception('load_members')
-
-    def put(self, path, **data):
-        if not self._base_uri:
-            self.load_members()
-        old_base_uri = self._base_uri
-        response = None
-        while response is None:
-            uri = self.client_url(path)
-            try:
-                logger.info('PUT %s', uri)
-                response = requests.put(uri, timeout=self.timeout, data=data)
-            except RequestException:
-                logger.exception('PUT %s data=%s', uri, data)
-                self._next_server()
-
-        if self._base_uri != old_base_uri:
-            try:
-                self.load_members()
-            except EtcdError:
-                logger.exception('load_members')
-        if response.status_code in [200, 201, 202, 204]:
-            return True
-        logger.error('Unexpected response: %s %s', response.status_code, response.content)
-        return False
-
-    def delete(self, path):
-        if not self._base_uri:
-            self.load_members()
-        old_base_uri = self._base_uri
-        response = None
-        while response is None:
-            uri = self.client_url(path)
-            try:
-                logger.info('DELETE %s', uri)
-                response = requests.delete(uri, timeout=self.timeout)
-            except RequestException:
-                logger.exception('DELETE %s', uri)
-                self._next_server()
-
-        if self._base_uri != old_base_uri:
-            try:
-                self.load_members()
-            except EtcdError:
-                logger.exception('load_members')
-        if response.status_code in [200, 202, 204]:
-            return True
-        logger.error('Unexpected response: %s %s', response.status_code, response.content)
-        return False
+            return not func(*args, **kwargs) is None
+        except etcd.EtcdException:
+            return False
+    return wrapper
 
 
 class Etcd(AbstractDCS):
@@ -197,118 +142,73 @@ class Etcd(AbstractDCS):
         while not client:
             try:
                 client = Client(config)
-            except EtcdError:
+            except etcd.EtcdException:
                 logger.info('waiting on etcd')
                 sleep(5)
         return client
 
-    def client_path(self, path):
-        return '/keys' + super(Etcd, self).client_path(path)
-
-    def get_client_path(self, path):
-        return self.client.get(self.client_path(path))
-
-    def put_client_path(self, path, **data):
-        return self.client.put(self.client_path(path), **data)
-
-    def delete_client_path(self, path):
-        try:
-            return self.client.delete(self.client_path(path))
-        except EtcdConnectionFailed:
-            return False
-
-    @staticmethod
-    def find_node(node, key):
-        """
-        >>> Etcd.find_node({}, None)
-        >>> Etcd.find_node({'dir': True, 'nodes': [], 'key': '/test/'}, 'test')
-        """
-        if not node.get('dir', False):
-            return None
-        key = node['key'] + key
-        for n in node['nodes']:
-            if n['key'] == key:
-                return n
-        return None
-
     @staticmethod
     def member(node):
-        conn_url, api_url = parse_connection_string(node['value'])
-        expiration = node.get('expiration', None)
-        ttl = node.get('ttl', None)
-        return Member(node['modifiedIndex'], node['key'].split('/')[-1], conn_url, api_url, expiration, ttl)
+        conn_url, api_url = parse_connection_string(node.value)
+        return Member(node.modifiedIndex, os.path.basename(node.key), conn_url, api_url, node.expiration, node.ttl)
 
     def get_cluster(self):
         try:
-            response, status_code = self.get_client_path('?recursive=true')
-            if status_code == 200:
-                node = self.find_node(response['node'], '/initialize')
-                initialize = True if node else False
-                # get list of members
-                node = self.find_node(response['node'], '/members') or {'nodes': []}
-                members = [self.member(n) for n in node['nodes']]
+            result = self.client.read(self.client_path(''), recursive=True)
+            nodes = {os.path.relpath(node.key, result.key): node for node in result.leaves}
 
-                # get last leader operation
-                last_leader_operation = 0
-                node = self.find_node(response['node'], '/optime')
-                if node:
-                    node = self.find_node(node, '/leader')
-                    if node:
-                        last_leader_operation = int(node['value'])
+            # get initialize flag
+            initialize = bool(nodes.get('initialize', False))
 
-                # get leader
-                leader = None
-                node = self.find_node(response['node'], '/leader')
-                if node:
-                    for m in members:
-                        if m.name == node['value']:
-                            leader = m
-                            break
-                    if not leader:
-                        leader = Member(-1, node['value'], None, None, None, None)
+            # get last leader operation
+            last_leader_operation = nodes.get('optime/leader', None)
+            last_leader_operation = 0 if last_leader_operation is None else int(last_leader_operation.value)
 
-                return Cluster(initialize, leader, last_leader_operation, members)
-            elif status_code == 404:
-                return Cluster(False, None, None, [])
+            # get list of members
+            members = [self.member(n) for k, n in nodes.items() if k.startswith('members/') and len(k.split('/')) == 2]
+
+            # get leader
+            leader = nodes.get('leader', None)
+            if leader:
+                leader = Member(-1, leader.value, None, None, None, None)
+                leader = ([m for m in members if m.name == leader.name] or [leader])[0]
+
+            return Cluster(initialize, leader, last_leader_operation, members)
+        except etcd.EtcdKeyNotFound:
+            return Cluster(False, None, None, [])
         except:
             logger.exception('get_cluster')
 
         raise EtcdError('Etcd is not responding properly')
 
+    @catch_etcd_errors
     def touch_member(self, connection_string, ttl=None):
-        try:
-            return self.put_client_path('/members/' + self._name, value=connection_string, ttl=ttl or self.member_ttl)
-        except EtcdError:
-            return False
+        return self.client.set(self.client_path('/members/' + self._name), connection_string, ttl or self.member_ttl)
 
+    @catch_etcd_errors
     def take_leader(self):
-        try:
-            return self.put_client_path('/leader', value=self._name, ttl=self.ttl)
-        except EtcdError:
-            return False
+        return self.client.set(self.client_path('/leader'), self._name, self.ttl)
 
+    @catch_etcd_errors
     def attempt_to_acquire_leader(self):
-        try:
-            ret = self.put_client_path('/leader', value=self._name, ttl=self.ttl, prevExist=False)
-            ret or logger.info('Could not take out TTL lock')
-            return ret
-        except EtcdError:
-            return False
+        ret = self.client.write(self.client_path('/leader'), self._name, ttl=self.ttl, prevExist=False)
+        ret or logger.info('Could not take out TTL lock')
+        return ret
 
+    @catch_etcd_errors
+    def write_leader_optime(self, state_handler):
+        return self.client.set(self.client_path('/optime/leader'), state_handler.last_operation())
+
+    @catch_etcd_errors
     def update_leader(self, state_handler):
-        if self.put_client_path('/leader', value=state_handler.name, ttl=self.ttl, prevValue=state_handler.name):
-            try:
-                self.put_client_path('/optime/leader', value=state_handler.last_operation())
-            except EtcdError:
-                pass
-            return True
-        return False
+        ret = self.client.test_and_set(self.client_path('/leader'), self._name, self._name, self.ttl)
+        ret and self.write_leader_optime(state_handler)
+        return ret
 
+    @catch_etcd_errors
     def race(self, path):
-        try:
-            return self.put_client_path(path, value=self._name, prevExist=False)
-        except EtcdError:
-            return False
+        return self.client.write(self.client_path(path), self._name, prevExist=False)
 
+    @catch_etcd_errors
     def delete_leader(self):
-        return self.delete_client_path('/leader?prevValue=' + self._name)
+        return self.client.delete(self.client_path('/leader'), prevValue=self._name)

patroni.py

@@ -12,6 +12,8 @@ from helpers.postgresql import Postgresql
 from helpers.utils import setup_signal_handlers, sleep
 from helpers.zookeeper import ZooKeeper
 
+logger = logging.getLogger(__name__)
+
 
 class Patroni:
 
@@ -44,7 +46,7 @@ class Patroni:
     def initialize(self):
         # wait for etcd to be available
         while not self.touch_member():
-            logging.info('waiting on DCS')
+            logger.info('waiting on DCS')
             sleep(5)
 
         # is data directory empty?
@@ -82,14 +84,14 @@ class Patroni:
 
         while True:
             self.touch_member()
-            logging.info(self.ha.run_cycle())
+            logger.info(self.ha.run_cycle())
             try:
                 if self.ha.state_handler.is_leader():
                     self.ha.cluster and self.ha.state_handler.create_replication_slots(self.ha.cluster)
                 else:
                     self.ha.state_handler.drop_replication_slots()
             except:
-                logging.exception('Exception when changing replication slots')
+                logger.exception('Exception when changing replication slots')
             self.schedule_next_run()
 
 

requirements-py2.txt

@@ -6,3 +6,4 @@ PyYAML
 requests
 six >= 1.7
 kazoo>=2.2.1
+python-etcd>=0.4.1

requirements-py3.txt

@@ -6,3 +6,4 @@ PyYAML
 requests
 six
 kazoo>=2.2.1
+python-etcd>=0.4.1

tests/test_api.py

@@ -1,10 +1,9 @@
 import psycopg2
-import sys
 import unittest
 
 from helpers.api import RestApiHandler, RestApiServer
-from test_postgresql import psycopg2_connect
 from six import BytesIO as IO
+from test_postgresql import psycopg2_connect
 
 
 def throws(*args, **kwargs):

tests/test_etcd.py

@@ -1,5 +1,6 @@
 import datetime
 import dns.resolver
+import etcd
 import json
 import requests
 import socket
@@ -7,8 +8,9 @@ import time
 import unittest
 
 from dns.exception import DNSException
-from helpers.dcs import Cluster, Member
-from helpers.etcd import Client, Etcd, EtcdConnectionFailed, EtcdError
+from helpers.dcs import Cluster, DCSError, Member
+from helpers.etcd import Client, Etcd
+from mock import Mock, patch
 
 
 class MockResponse:
@@ -21,6 +23,17 @@ class MockResponse:
     def json(self):
         return json.loads(self.content)
 
+    @property
+    def data(self):
+        return self.content
+
+    @property
+    def status(self):
+        return self.status_code
+
+    def getheader(*args):
+        return ''
+
 
 class MockPostgresql:
     name = ''
@@ -30,51 +43,64 @@ class MockPostgresql:
 
 
 def requests_get(url, **kwargs):
-    members = '[{"id":14855829450254237642,"peerURLs":["http://localhost:2380","http://localhost:7001"],"name":"default","clientURLs":["http://localhost:2379","http://localhost:4001"]}]'
+    members = '[{"id":14855829450254237642,"peerURLs":["http://localhost:2380","http://localhost:7001"],' +\
+              '"name":"default","clientURLs":["http://localhost:2379","http://localhost:4001"]}]'
     response = MockResponse()
-    if url.endswith('/v2/members'):
-        response.content = '{"members": ' + members + '}'
-        if url.startswith('http://error'):
-            response.status_code = 404
+    if url.startswith('http://local'):
+        raise requests.exceptions.RequestException()
     elif url.endswith('/members'):
         if url.startswith('http://error'):
             response.content = '[{}]'
         else:
             response.content = members
-    elif url.endswith('/bad_response'):
-        response.content = '{'
     elif url.startswith('http://exhibitor'):
         response.content = '{"servers":["127.0.0.1","127.0.0.2","127.0.0.3"],"port":2181}'
-    elif url.startswith('http://local'):
-        raise requests.exceptions.RequestException()
-    elif url.startswith('http://remote') or url.startswith('http://127.0.0.1') or url.startswith('http://error'):
-        response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql1","value":"postgres://replicator:[email protected]:5434/postgres?application_name=http://127.0.0.1:8009/patroni","expiration":"2015-05-15T09:10:59.949384522Z","ttl":21,"modifiedIndex":20727,"createdIndex":20727},{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:[email protected]:5433/postgres?application_name=http://127.0.0.1:8008/patroni","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}'
-    elif url.startswith('http://other'):
-        response.status_code = 404
-    elif url.startswith('http://noleader'):
-        response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:[email protected]:5433/postgres?application_name=http://127.0.0.1:8008/patroni","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}'
     else:
         response.status_code = 404
         response.ok = False
     return response
 
 
-def requests_put(url, **kwargs):
-    if url.startswith('http://local') or '/optime/leader' in url:
-        raise requests.exceptions.RequestException()
-    response = MockResponse()
-    response.status_code = 201
-    if url.startswith('http://other'):
-        response.status_code = 404
-    return response
+def etcd_write(key, value, **kwargs):
+    if key == '/service/test/leader':
+        if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True):
+            return True
+    raise etcd.EtcdException
 
 
-def requests_delete(url, **kwargs):
-    if url.startswith('http://local'):
-        raise requests.exceptions.RequestException()
-    response = MockResponse()
-    response.status_code = 503 if url.startswith('http://error') else 204
-    return response
+def etcd_delete(key, **kwargs):
+    raise etcd.EtcdException
+
+
+def etcd_read(key, **kwargs):
+    if key == '/service/noleader':
+        raise DCSError('noleader')
+    elif key == '/service/nocluster':
+        raise etcd.EtcdKeyNotFound
+
+    response = {"action": "get", "node": {"key": "/service/batman5", "dir": True, "nodes": [
+                {"key": "/service/batman5/initialize", "value": "postgresql0",
+                 "modifiedIndex": 1582, "createdIndex": 1582},
+                {"key": "/service/batman5/leader", "value": "postgresql1",
+                 "expiration": "2015-05-15T09:11:00.037397538Z", "ttl": 21,
+                 "modifiedIndex": 20728, "createdIndex": 20434},
+                {"key": "/service/batman5/optime", "dir": True, "nodes": [
+                    {"key": "/service/batman5/optime/leader", "value": "2164261704",
+                     "modifiedIndex": 20729, "createdIndex": 20729}],
+                 "modifiedIndex": 20437, "createdIndex": 20437},
+                {"key": "/service/batman5/members", "dir": True, "nodes": [
+                    {"key": "/service/batman5/members/postgresql1",
+                     "value": "postgres://replicator:[email protected]:5434/postgres"
+                        + "?application_name=http://127.0.0.1:8009/patroni",
+                     "expiration": "2015-05-15T09:10:59.949384522Z", "ttl": 21,
+                     "modifiedIndex": 20727, "createdIndex": 20727},
+                    {"key": "/service/batman5/members/postgresql0",
+                     "value": "postgres://replicator:[email protected]:5433/postgres"
+                        + "?application_name=http://127.0.0.1:8008/patroni",
+                     "expiration": "2015-05-15T09:11:09.611860899Z", "ttl": 30,
+                     "modifiedIndex": 20730, "createdIndex": 20730}],
+                 "modifiedIndex": 1581, "createdIndex": 1581}], "modifiedIndex": 1581, "createdIndex": 1581}}
+    return etcd.EtcdResult(**response)
 
 
 def time_sleep(_):
@@ -104,6 +130,12 @@ def socket_getaddrinfo(*args):
     raise socket.error()
 
 
+def http_request(method, url, **kwargs):
+    if url == 'http://localhost:2379/':
+        return MockResponse()
+    raise socket.error
+
+
 class TestMember(unittest.TestCase):
 
     def __init__(self, method_name='runTest'):
@@ -125,49 +157,33 @@ class TestClient(unittest.TestCase):
     def set_up(self):
         socket.getaddrinfo = socket_getaddrinfo
         requests.get = requests_get
-        requests.put = requests_put
-        requests.delete = requests_delete
         dns.resolver.query = dns_query
-        self.client = Client({'discovery_srv': 'test'})
+        with patch.object(etcd.Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
+            self.client = Client({'discovery_srv': 'test'})
+            self.client.http.request = http_request
 
-    def test__get(self):
-        self.assertRaises(EtcdError, self.client._get, '/bad_response')
+    def test_api_execute(self):
+        self.client._base_uri = 'http://localhost:4001'
+        self.client._machines_cache = ['http://localhost:2379']
+        self.client.api_execute('/', 'GET')
 
     def test_get_srv_record(self):
-        self.assertEquals(Client.get_srv_record('blabla'), [])
-        self.assertEquals(Client.get_srv_record('exception'), [])
+        self.assertEquals(self.client.get_srv_record('blabla'), [])
+        self.assertEquals(self.client.get_srv_record('exception'), [])
 
-    def test_get_client_urls_from_dns(self):
-        self.assertEquals(Client.get_client_urls_from_dns('ok:2379'), ['http://127.0.0.1:2379/v2'])
+    def test__get_machines_cache_from_srv(self):
+        self.client.get_srv_record = lambda e: [('localhost', 2380)]
+        self.client._get_machines_cache_from_srv('blabla')
 
-    def test_load_members(self):
-        self.client._base_uri = self.client._base_uri.replace('localhost', 'error_code')
-        self.assertRaises(EtcdError, self.client.load_members)
-        self.client._base_uri = 'http://error_code:2380'
-        self.assertRaises(EtcdError, self.client.load_members)
-        self.client._base_uri = None
+    def test__get_machines_cache_from_dns(self):
+        self.client._get_machines_cache_from_dns('ok:2379')
+
+    def test__load_machines_cache(self):
         self.client._config = {}
-        self.assertRaises(Exception, self.client.load_members)
-
-    def test_get(self):
-        self.client._base_uri = None
-        self.assertRaises(EtcdConnectionFailed, self.client.get, '')
-        self.client._members_cache = ['http://error_code:4001/v2']
-        self.client.get('')
-
-    def test_put(self):
-        self.client._base_uri = None
-        self.assertRaises(EtcdConnectionFailed, self.client.put, '')
-        self.client._base_uri = 'http://localhost:4001/v2'
-        self.client._members_cache = ['http://error_code:4001/v2']
-        self.client.put('')
-
-    def test_delete(self):
-        self.client._base_uri = None
-        self.assertRaises(EtcdConnectionFailed, self.client.delete, '')
-        self.client._base_uri = 'http://localhost:4001/v2'
-        self.client._members_cache = ['http://error_code:4001/v2']
-        self.client.delete('')
+        self.assertRaises(Exception, self.client._load_machines_cache)
+        self.client._config = {'discovery_srv': 'blabla'}
+        self.assertRaises(etcd.EtcdException, self.client._load_machines_cache)
 
 
 class TestEtcd(unittest.TestCase):
@@ -177,37 +193,29 @@ class TestEtcd(unittest.TestCase):
         super(TestEtcd, self).__init__(method_name)
 
     def set_up(self):
-        socket.getaddrinfo = socket_getaddrinfo
-        requests.get = requests_get
-        requests.put = requests_put
-        requests.delete = requests_delete
         time.sleep = time_sleep
-        self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
+        with patch.object(Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
+            self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
+            self.etcd.client.write = etcd_write
+            self.etcd.client.read = etcd_read
 
     def test_get_etcd_client(self):
         time.sleep = time_sleep_exception
-        self.assertRaises(Exception, self.etcd.get_etcd_client, {'host': 'error:2379'})
-
-    def test_get_client_path(self):
-        self.assertRaises(Exception, self.etcd.get_client_path, '', 2)
-
-    def test_put_client_path(self):
-        self.assertRaises(EtcdError, self.etcd.put_client_path, '')
-
-    def test_delete_client_path(self):
-        self.assertFalse(self.etcd.delete_client_path(''))
+        with patch.object(etcd.Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(side_effect=etcd.EtcdException)
+            self.assertRaises(Exception, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
 
     def test_get_cluster(self):
-        self.assertRaises(EtcdError, self.etcd.get_cluster)
-        self.etcd.client._base_uri = self.etcd.client._base_uri.replace('local', 'remote')
+        self.assertIsInstance(self.etcd.get_cluster(), Cluster)
+        self.etcd._base_path = '/service/nocluster'
         cluster = self.etcd.get_cluster()
         self.assertIsInstance(cluster, Cluster)
-        self.etcd.client._base_uri = self.etcd.client._base_uri.replace('remote', 'other')
-        self.etcd.get_cluster()
-        self.etcd.client._base_uri = self.etcd.client._base_uri.replace('other', 'noleader')
-        self.etcd.get_cluster()
+        self.assertIsNone(cluster.leader)
 
     def test_current_leader(self):
+        self.assertIsInstance(self.etcd.current_leader(), Member)
+        self.etcd._base_path = '/service/noleader'
         self.assertIsNone(self.etcd.current_leader())
 
     def test_touch_member(self):
@@ -216,14 +224,12 @@ class TestEtcd(unittest.TestCase):
     def test_take_leader(self):
         self.assertFalse(self.etcd.take_leader())
 
-    def test_attempt_to_acquire_leader(self):
-        self.assertFalse(self.etcd.attempt_to_acquire_leader())
-
     def test_update_leader(self):
-        url = self.etcd.client._base_uri = self.etcd.client._base_uri.replace('local', 'remote')
         self.assertTrue(self.etcd.update_leader(MockPostgresql()))
-        self.etcd.client._base_uri = url.replace('remote', 'other')
-        self.assertFalse(self.etcd.update_leader(MockPostgresql()))
 
     def test_race(self):
         self.assertFalse(self.etcd.race(''))
+
+    def test_delete_leader(self):
+        self.etcd.client.delete = etcd_delete
+        self.assertFalse(self.etcd.delete_leader())

tests/test_ha.py

@@ -1,10 +1,10 @@
 import unittest
-import requests
 
-from helpers.dcs import DCSError
-from helpers.etcd import Cluster, Etcd
+from helpers.dcs import Cluster, DCSError
+from helpers.etcd import Client, Etcd
 from helpers.ha import Ha
-from test_etcd import requests_get, requests_put, requests_delete
+from mock import Mock, patch
+from test_etcd import etcd_read, etcd_write
 
 
 def true(*args, **kwargs):
@@ -71,15 +71,16 @@ class TestHa(unittest.TestCase):
         super(TestHa, self).__init__(method_name)
 
     def set_up(self):
-        requests.get = requests_get
-        requests.put = requests_put
-        requests.delete = requests_delete
         self.p = MockPostgresql()
-        self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
-        self.ha = Ha(self.p, self.e)
-        self.ha.load_cluster_from_dcs()
-        self.ha.cluster = get_unlocked_cluster()
-        self.ha.load_cluster_from_dcs = nop
+        with patch.object(Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+            self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
+            self.e.client.read = etcd_read
+            self.e.client.write = etcd_write
+            self.ha = Ha(self.p, self.e)
+            self.ha.load_cluster_from_dcs()
+            self.ha.cluster = get_unlocked_cluster()
+            self.ha.load_cluster_from_dcs = nop
 
     def test_load_cluster_from_dcs(self):
         ha = Ha(self.p, self.e)

tests/test_patroni.py

@@ -1,18 +1,20 @@
 import datetime
 import helpers.zookeeper
 import psycopg2
-import requests
 import subprocess
 import sys
 import time
 import unittest
 import yaml
 
-from patroni import Patroni, main
+from helpers.api import RestApiServer
 from helpers.dcs import Cluster, Member
+from helpers.etcd import Etcd
 from helpers.zookeeper import ZooKeeper
+from mock import Mock, patch
+from patroni import Patroni, main
 from six.moves import BaseHTTPServer
-from test_etcd import requests_get, requests_put, requests_delete
+from test_etcd import Client, etcd_read, etcd_write
 from test_ha import true, false
 from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
 from test_zookeeper import MockKazooClient
@@ -26,6 +28,15 @@ def time_sleep(*args):
     raise Exception()
 
 
+class Mock_BaseServer__is_shut_down:
+
+    def set(self):
+        pass
+
+    def clear(self):
+        pass
+
+
 class TestPatroni(unittest.TestCase):
 
     def __init__(self, method_name='runTest'):
@@ -37,9 +48,6 @@ class TestPatroni(unittest.TestCase):
         self.touched = False
         subprocess.call = subprocess_call
         psycopg2.connect = psycopg2_connect
-        requests.get = requests_get
-        requests.put = requests_put
-        requests.delete = requests_delete
         self.time_sleep = time.sleep
         time.sleep = nop
         self.write_pg_hba = Postgresql.write_pg_hba
@@ -47,9 +55,14 @@ class TestPatroni(unittest.TestCase):
         Postgresql.write_pg_hba = nop
         Postgresql.write_recovery_conf = nop
         BaseHTTPServer.HTTPServer.__init__ = nop
+        RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down()
+        RestApiServer._BaseServer__shutdown_request = True
+        RestApiServer.socket = 0
         with open('postgres0.yml', 'r') as f:
             config = yaml.load(f)
-            self.g = Patroni(config)
+            with patch.object(Client, 'machines') as mock_machines:
+                mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+                self.p = Patroni(config)
 
     def tear_down(self):
         time.sleep = self.time_sleep
@@ -58,50 +71,74 @@ class TestPatroni(unittest.TestCase):
 
     def test_get_dcs(self):
         helpers.zookeeper.KazooClient = MockKazooClient
-        self.assertIsInstance(self.g.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
-        self.assertRaises(Exception, self.g.get_dcs, '', {})
+        self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
+        self.assertRaises(Exception, self.p.get_dcs, '', {})
 
     def test_patroni_main(self):
         main()
         sys.argv = ['patroni.py', 'postgres0.yml']
         time.sleep = time_sleep
-        self.assertRaises(Exception, main)
+
+        with patch.object(Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+            Patroni.initialize = nop
+            touch_member = Patroni.touch_member
+            run = Patroni.run
+
+            Patroni.touch_member = self.touch_member
+            Patroni.run = time_sleep
+
+            Etcd.delete_leader = nop
+
+            self.assertRaises(Exception, main)
+
+            Patroni.run = run
+            Patroni.touch_member = touch_member
 
     def test_patroni_run(self):
         time.sleep = time_sleep
-        self.g.postgresql.is_leader = lambda: False
-        self.g.ha.state_handler.sync_replication_slots = time_sleep
-        self.assertRaises(Exception, self.g.run)
+        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.assertRaises(Exception, self.p.run)
+        self.p.ha.state_handler.is_leader = lambda: False
+        self.p.api.start = nop
+        self.assertRaises(Exception, self.p.run)
 
-    def touch_member(self):
+    def touch_member(self, ttl=None):
         if not self.touched:
             self.touched = True
             return False
         return True
 
     def test_touch_member(self):
+        self.p.ha.dcs.client.write = etcd_write
+        self.p.touch_member()
         now = datetime.datetime.utcnow()
-        member = Member(0, self.g.postgresql.name, 'b', 'c', (now + datetime.timedelta(
-            seconds=self.g.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
-        self.g.ha.cluster = Cluster(True, member, 0, [member])
-        self.g.touch_member()
+        member = Member(0, self.p.postgresql.name, 'b', 'c', (now + datetime.timedelta(
+            seconds=self.p.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
+        self.p.ha.cluster = Cluster(True, member, 0, [member])
+        self.p.touch_member()
 
     def test_patroni_initialize(self):
-        self.g.postgresql.should_use_s3_to_create_replica = false
-        self.g.ha.dcs.client._base_uri = 'http://remote'
-        self.g.postgresql.data_directory_empty = true
-        self.g.ha.dcs.race = true
-        self.g.initialize()
-        self.g.ha.dcs.race = false
-        self.g.initialize()
-        self.g.postgresql.data_directory_empty = false
-        self.g.touch_member = self.touch_member
-        self.g.initialize()
-        self.g.postgresql.data_directory_empty = true
+        self.p.postgresql.should_use_s3_to_create_replica = false
+        self.p.ha.dcs.client.write = etcd_write
+        self.p.touch_member = self.touch_member
+        self.p.postgresql.data_directory_empty = true
+        self.p.ha.dcs.race = true
+        self.p.initialize()
+
+        self.p.ha.dcs.race = false
         time.sleep = time_sleep
-        self.g.postgresql.sync_from_leader = false
-        self.assertRaises(Exception, self.g.initialize)
+        self.p.ha.dcs.client.read = etcd_read
+        self.p.initialize()
+
+        self.p.ha.dcs.current_leader = nop
+        self.assertRaises(Exception, self.p.initialize)
+
+        self.p.postgresql.data_directory_empty = false
+        self.p.initialize()
 
     def test_schedule_next_run(self):
-        self.g.next_run = time.time() - self.g.nap_time - 1
-        self.g.schedule_next_run()
+        self.p.next_run = time.time() - self.p.nap_time - 1
+        self.p.schedule_next_run()

tests/test_postgresql.py

@@ -1,6 +1,3 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
 import os
 import psycopg2
 import shutil

tests/test_zookeeper.py

@@ -1,6 +1,6 @@
 import helpers.zookeeper
-import unittest
 import requests
+import unittest
 
 from helpers.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError
 from kazoo.client import KazooState