Merge pull request #45 from zalando/feature/query-retry

Refactor Postgresql.query method to use common retry mechanism

patroni/api.py

@@ -4,6 +4,8 @@ import json
 import logging
 import psycopg2
 
+from patroni.exceptions import PostgresConnectionException
+from patroni.utils import Retry, RetryFailedError
 from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 from six.moves.socketserver import ThreadingMixIn
 from threading import Thread
@@ -59,6 +61,14 @@ class RestApiHandler(BaseHTTPRequestHandler):
         self.end_headers()
         self.wfile.write(b'Hello!')
 
+    def do_GET_patroni(self):
+        response = self.get_postgresql_status(True)
+
+        self.send_response(200)
+        self.send_header('Content-Type', 'application/json')
+        self.end_headers()
+        self.wfile.write(json.dumps(response).encode('utf-8'))
+
     def parse_request(self):
         """Override parse_request method to enrich basic functionality of `BaseHTTPRequestHandler` class
 
@@ -77,16 +87,22 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 self.command = mname
         return ret
 
-    def get_postgresql_status(self):
+    def query(self, sql, *params, **kwargs):
+        if not kwargs.get('retry', False):
+            return self.server.query(sql, *params)
+        retry = Retry(delay=2, retry_exceptions=PostgresConnectionException)
+        return retry(self.server.query, sql, *params)
+
+    def get_postgresql_status(self, retry=False):
         try:
-            row = self.server.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
-                                              pg_is_in_recovery(),
-                                              CASE WHEN pg_is_in_recovery()
-                                                   THEN null
-                                                   ELSE pg_current_xlog_location() END,
-                                              pg_last_xlog_receive_location(),
-                                              pg_last_xlog_replay_location(),
-                                              pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")[0]
+            row = self.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+                                       pg_is_in_recovery(),
+                                       CASE WHEN pg_is_in_recovery()
+                                            THEN null
+                                            ELSE pg_current_xlog_location() END,
+                                       pg_last_xlog_receive_location(),
+                                       pg_last_xlog_replay_location(),
+                                       pg_is_in_recovery() AND pg_is_xlog_replay_paused()""", retry=retry)[0]
             return {
                 'running': True,
                 'postmaster_start_time': row[0],
@@ -98,7 +114,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
                     'location': row[2]
                 })
             }
-        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+        except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
             logger.exception('get_postgresql_status')
             return {'running': self.server.patroni.postgresql.is_running()}
 
@@ -128,11 +144,15 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
         self.daemon = True
 
     def query(self, sql, *params):
-        cursor = self.patroni.postgresql.connection().cursor()
-        cursor.execute(sql, params)
-        ret = [r for r in cursor]
-        cursor.close()
-        return ret
+        cursor = None
+        try:
+            with self.patroni.postgresql.connection().cursor() as cursor:
+                cursor.execute(sql, params)
+                return [r for r in cursor]
+        except psycopg2.Error as e:
+            if cursor and cursor.connection.closed == 0:
+                raise e
+            raise PostgresConnectionException('connection problems')
 
     @staticmethod
     def _set_fd_cloexec(fd):

patroni/exceptions.py

@@ -7,7 +7,7 @@ class PatroniException(Exception):
 
     def __str__(self):
         """
-        >>> str(DCSError('foo'))
+        >>> str(PatroniException('foo'))
         "'foo'"
         """
         return repr(self.value)
@@ -19,3 +19,7 @@ class PostgresException(PatroniException):
 
 class DCSError(PatroniException):
     pass
+
+
+class PostgresConnectionException(PostgresException):
+    pass

patroni/ha.py

@@ -1,7 +1,7 @@
 import logging
 import psycopg2
 
-from patroni.dcs import DCSError
+from patroni.exceptions import DCSError, PostgresConnectionException
 
 logger = logging.getLogger(__name__)
 
@@ -97,5 +97,5 @@ class Ha:
             if self.state_handler.is_leader():
                 self.state_handler.demote(None)
                 return 'demoted self because DCS is not accessible and i was a leader'
-        except psycopg2.Error:
+        except (psycopg2.Error, PostgresConnectionException):
             logger.exception('Error communicating with Postgresql.  Will try again')

patroni/postgresql.py

@@ -6,8 +6,8 @@ import shutil
 import subprocess
 import time
 
-from patroni.exceptions import PostgresException
-from patroni.utils import sleep
+from patroni.exceptions import PostgresConnectionException, PostgresException
+from patroni.utils import Retry, RetryFailedError
 from six.moves.urllib_parse import urlparse
 
 logger = logging.getLogger(__name__)
@@ -68,6 +68,7 @@ class Postgresql:
         self._connection = None
         self._cursor_holder = None
         self.members = []  # list of already existing replication slots
+        self.retry = Retry(max_tries=-1, deadline=10, max_delay=1, retry_exceptions=PostgresConnectionException)
 
     def get_local_address(self):
         listen_addresses = self.listen_addresses.split(',')
@@ -87,32 +88,26 @@ class Postgresql:
         return self._connection
 
     def _cursor(self):
-        if not self._cursor_holder or self._cursor_holder.closed:
+        if not self._cursor_holder or self._cursor_holder.closed or self._cursor_holder.connection.closed != 0:
             self._cursor_holder = self.connection().cursor()
         return self._cursor_holder
 
-    def disconnect(self):
-        self._connection and self._connection.close()
-        self._connection = self._cursor_holder = None
+    def _query(self, sql, *params):
+        cursor = None
+        try:
+            cursor = self._cursor()
+            cursor.execute(sql, params)
+            return cursor
+        except psycopg2.Error as e:
+            if cursor and cursor.connection.closed == 0:
+                raise e
+            raise PostgresConnectionException('connection problems')
 
     def query(self, sql, *params):
-        max_attempts = 0
-        while True:
-            ex = None
-            try:
-                cursor = self._cursor()
-                cursor.execute(sql, params)
-                return cursor
-            except psycopg2.Error as e:
-                if self._connection and self._connection.closed == 0:
-                    raise e
-                ex = e
-            if ex:
-                self.disconnect()
-                max_attempts += 1
-                if max_attempts >= 3:
-                    raise ex
-                sleep(5)
+        try:
+            return self.retry(self._query, sql, *params)
+        except RetryFailedError as e:
+            raise PostgresConnectionException(str(e))
 
     def data_directory_empty(self):
         return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []

patroni/utils.py

@@ -6,7 +6,7 @@ import signal
 import sys
 import time
 
-from patroni.exceptions import DCSError
+from patroni.exceptions import PatroniException
 
 ignore_sigterm = False
 interrupted_sleep = False
@@ -90,7 +90,7 @@ def reap_children():
             reap_children = False
 
 
-class RetryFailedError(DCSError):
+class RetryFailedError(PatroniException):
 
     """Raised when retrying an operation ultimately failed, after retrying the maximum number of attempts."""
 
@@ -100,7 +100,7 @@ class Retry:
     """Helper for retrying a method in the face of retry-able exceptions"""
 
     def __init__(self, max_tries=1, delay=0.1, backoff=2, max_jitter=0.8, max_delay=3600,
-                 sleep_func=time.sleep, deadline=None, retry_exceptions=DCSError):
+                 sleep_func=sleep, deadline=None, retry_exceptions=PatroniException):
         """Create a :class:`Retry` instance for retrying function calls
 
         :param max_tries: How many times to retry the command. -1 means infinite tries.
@@ -154,13 +154,10 @@ class Retry:
                 if self._attempts == self.max_tries:
                     raise RetryFailedError("Too many retry attempts")
                 self._attempts += 1
-                sleeptime = self._cur_delay + (
-                    random.randint(0, self.max_jitter) / 100.0)
+                sleeptime = self._cur_delay + (random.randint(0, self.max_jitter) / 100.0)
 
-                if self._cur_stoptime is not None and \
-                   time.time() + sleeptime >= self._cur_stoptime:
+                if self._cur_stoptime is not None and time.time() + sleeptime >= self._cur_stoptime:
                     raise RetryFailedError("Exceeded retry deadline")
                 else:
                     self.sleep_func(sleeptime)
-                self._cur_delay = min(self._cur_delay * self.backoff,
-                                      self.max_delay)
+                self._cur_delay = min(self._cur_delay * self.backoff, self.max_delay)

tests/test_api.py

@@ -1,35 +1,14 @@
 import psycopg2
 import unittest
-import ssl
 
+from mock import Mock, patch
 from patroni.api import RestApiHandler, RestApiServer
 from six import BytesIO as IO
 from six.moves import BaseHTTPServer
-from test_postgresql import psycopg2_connect
+from test_postgresql import psycopg2_connect, MockCursor
 
 
-def nop(*args, **kwargs):
-    pass
-
-
-def throws(*args, **kwargs):
-    raise psycopg2.OperationalError()
-
-
-def ssl_wrap_socket(socket, *args, **kwargs):
-    return socket
-
-
-class Mock_BaseServer__is_shut_down:
-
-    def set(self):
-        pass
-
-    def clear(self):
-        pass
-
-
-class MockPostgresql:
+class MockPostgresql(Mock):
 
     def connection(self):
         return psycopg2_connect()
@@ -40,8 +19,7 @@ class MockPostgresql:
 
 class MockPatroni:
 
-    def __init__(self):
-        self.postgresql = MockPostgresql()
+    postgresql = MockPostgresql()
 
 
 class MockRequest:
@@ -55,32 +33,35 @@ class MockRequest:
 
 class MockRestApiServer(RestApiServer):
 
-    def __init__(self, Handler, path, *args):
+    def __init__(self, Handler, path):
+        self.socket = 0
+        BaseHTTPServer.HTTPServer.__init__ = Mock()
+        MockRestApiServer._BaseServer__is_shut_down = Mock()
+        MockRestApiServer._BaseServer__shutdown_request = True
         config = {'listen': '127.0.0.1:8008', 'auth': 'test:test', 'certfile': 'dumb'}
         super(MockRestApiServer, self).__init__(MockPatroni(), config)
-        if len(args) > 0:
-            self.query = args[0]
         Handler(MockRequest(path), ('0.0.0.0', 8080), self)
 
 
+@patch('ssl.wrap_socket', Mock(return_value=0))
 class TestRestApiHandler(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestRestApiHandler, self).__init__(method_name)
-
-    def set_up(self):
-        BaseHTTPServer.HTTPServer.__init__ = nop
-        RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down()
-        RestApiServer._BaseServer__shutdown_request = True
-        RestApiServer.socket = 0
-        ssl.wrap_socket = ssl_wrap_socket
-
     def test_do_GET(self):
         MockRestApiServer(RestApiHandler, b'GET /')
-        MockRestApiServer(RestApiHandler, b'GET /', throws)
+        with patch.object(RestApiServer, 'query', Mock(side_effect=psycopg2.OperationalError())):
+            MockRestApiServer(RestApiHandler, b'GET /')
 
     def test_do_GET_sampleauth(self):
         MockRestApiServer(RestApiHandler, b'GET /sampleauth')
         MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization:')
         MockRestApiServer(RestApiHandler, b'GET /sampleauth\nAuthorization: Basic dGVzdDp0ZXN0')
+
+    def test_do_GET_patroni(self):
+        MockRestApiServer(RestApiHandler, b'GET /patroni')
+
+    @patch('time.sleep', Mock())
+    def test_RestApiServer_query(self):
+        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')
+        with patch.object(MockPostgresql, 'connection', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')

tests/test_etcd.py

@@ -1,11 +1,9 @@
 import datetime
-import dns.resolver
 import etcd
 import json
 import requests
 import urllib3
 import socket
-import time
 import unittest
 
 from dns.exception import DNSException
@@ -40,8 +38,7 @@ class MockResponse:
         return ''
 
 
-class MockPostgresql:
-    name = ''
+class MockPostgresql(Mock):
 
     def last_operation(self):
         return '0'
@@ -88,10 +85,6 @@ def etcd_write(key, value, **kwargs):
     raise etcd.EtcdException
 
 
-def etcd_delete(key, **kwargs):
-    raise etcd.EtcdException
-
-
 def etcd_read(key, **kwargs):
     if key == '/service/noleader/':
         raise DCSError('noleader')
@@ -123,18 +116,10 @@ def etcd_read(key, **kwargs):
     return etcd.EtcdResult(**response)
 
 
-def time_sleep(_):
-    pass
-
-
 class SleepException(Exception):
     pass
 
 
-def time_sleep_exception(_):
-    raise SleepException()
-
-
 class MockSRV:
     port = 2380
     target = '127.0.0.1'
@@ -151,7 +136,7 @@ def dns_query(name, type):
 def socket_getaddrinfo(*args):
     if args[0] == 'ok':
         return [(2, 1, 6, '', ('127.0.0.1', 2379)), (2, 1, 6, '', ('127.0.0.1', 2379))]
-    raise socket.error()
+    raise socket.error
 
 
 def http_request(method, url, **kwargs):
@@ -162,9 +147,6 @@ def http_request(method, url, **kwargs):
 
 class TestMember(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        super(TestMember, self).__init__(method_name)
-
     def test_real_ttl(self):
         now = datetime.datetime.utcnow()
         member = Member(0, 'a', 'b', 'c', (now + datetime.timedelta(seconds=2)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
@@ -172,16 +154,14 @@ class TestMember(unittest.TestCase):
         self.assertEquals(Member(0, 'a', 'b', 'c', '', None).real_ttl(), -1)
 
 
+@patch('dns.resolver.query', dns_query)
+@patch('socket.getaddrinfo', socket_getaddrinfo)
+@patch('requests.get', requests_get)
 class TestClient(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestClient, self).__init__(method_name)
-
-    def set_up(self):
-        socket.getaddrinfo = socket_getaddrinfo
-        requests.get = requests_get
-        dns.resolver.query = dns_query
+    @patch('dns.resolver.query', dns_query)
+    @patch('requests.get', requests_get)
+    def setUp(self):
         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'})
@@ -206,11 +186,11 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._result_from_response, response)
 
     def test__get_machines_cache_from_srv(self):
-        self.client.get_srv_record = lambda e: [('localhost', 2380)]
+        self.client.get_srv_record = Mock(return_value=[('localhost', 2380)])
         self.client._get_machines_cache_from_srv('blabla')
 
     def test__get_machines_cache_from_dns(self):
-        self.client._get_machines_cache_from_dns('ok:2379')
+        self.client._get_machines_cache_from_dns('error:2379')
 
     def test__load_machines_cache(self):
         self.client._config = {}
@@ -219,25 +199,24 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._load_machines_cache)
 
 
+@patch('time.sleep', Mock())
+@patch('requests.get', requests_get)
 class TestEtcd(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestEtcd, self).__init__(method_name)
-
-    def set_up(self):
-        time.sleep = time_sleep
+    def setUp(self):
         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
+            self.etcd.client.delete = Mock(side_effect=etcd.EtcdException())
 
+    @patch('dns.resolver.query', dns_query)
     def test_get_etcd_client(self):
-        time.sleep = time_sleep_exception
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(side_effect=etcd.EtcdException)
-            self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
+            with patch('time.sleep', Mock(side_effect=SleepException())):
+                self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
 
     def test_get_cluster(self):
         self.assertIsInstance(self.etcd.get_cluster(), Cluster)
@@ -257,7 +236,7 @@ class TestEtcd(unittest.TestCase):
     def test_take_leader(self):
         self.assertFalse(self.etcd.take_leader())
 
-    def testattempt_to_acquire_leader(self):
+    def test_attempt_to_acquire_leader(self):
         self.etcd._base_path = '/service/exists'
         self.assertFalse(self.etcd.attempt_to_acquire_leader())
         self.etcd._base_path = '/service/failed'
@@ -270,11 +249,9 @@ class TestEtcd(unittest.TestCase):
         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())
 
     def test_watch(self):

tests/test_ha.py

@@ -4,7 +4,7 @@ from mock import Mock, patch
 from patroni.dcs import Cluster, DCSError
 from patroni.etcd import Client, Etcd
 from patroni.ha import Ha
-from test_etcd import etcd_read, etcd_write
+from test_etcd import socket_getaddrinfo, etcd_read, etcd_write
 
 
 def true(*args, **kwargs):
@@ -52,35 +52,24 @@ class MockPostgresql:
         return 0
 
 
-def nop(*args, **kwargs):
-    pass
-
-
-def dead_etcd():
-    raise DCSError('Etcd is not responding properly')
-
-
 def get_unlocked_cluster():
     return Cluster(False, None, None, [])
 
 
 class TestHa(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestHa, self).__init__(method_name)
-
-    def set_up(self):
+    @patch('socket.getaddrinfo', socket_getaddrinfo)
+    def setUp(self):
         self.p = MockPostgresql()
         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 = Etcd('foo', {'ttl': 30, 'host': 'ok: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
+            self.ha.load_cluster_from_dcs = Mock()
 
     def test_load_cluster_from_dcs(self):
         ha = Ha(self.p, self.e)
@@ -144,5 +133,5 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.run_cycle(), 'no action.  i am a secondary and i am following a leader')
 
     def test_no_etcd_connection_master_demote(self):
-        self.ha.load_cluster_from_dcs = dead_etcd
+        self.ha.load_cluster_from_dcs = Mock(side_effect=DCSError('Etcd is not responding properly'))
         self.assertEquals(self.ha.run_cycle(), 'demoted self because DCS is not accessible and i was a leader')

tests/test_patroni.py

@@ -1,7 +1,4 @@
 import datetime
-import patroni.zookeeper
-import psycopg2
-import subprocess
 import sys
 import time
 import unittest
@@ -15,41 +12,16 @@ from patroni.exceptions import DCSError, PostgresException
 from patroni import Patroni, main
 from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
-from test_api import Mock_BaseServer__is_shut_down
-from test_etcd import Client, etcd_read, etcd_write
+from test_etcd import Client, SleepException, etcd_read, etcd_write
 from test_ha import true, false
-from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
+from test_postgresql import Postgresql, psycopg2_connect
 from test_zookeeper import MockKazooClient
 
 
-def nop(*args, **kwargs):
-    pass
-
-
-class SleepException(Exception):
-    pass
-
-
 def time_sleep(*args):
     raise SleepException()
 
 
-def keyboard_interrupt(*args):
-    raise KeyboardInterrupt
-
-
-class Mock_BaseServer__is_shut_down:
-
-    def wait(self):
-        pass
-
-    def set(self):
-        pass
-
-    def clear(self):
-        pass
-
-
 def get_cluster(initialize, leader):
     return Cluster(initialize, leader, None, None)
 
@@ -78,26 +50,18 @@ def get_cluster_dcs_error():
     raise DCSError('')
 
 
+@patch('time.sleep', Mock())
+@patch('subprocess.call', Mock(return_value=0))
+@patch('psycopg2.connect', psycopg2_connect)
+@patch.object(Postgresql, 'write_pg_hba', Mock())
+@patch.object(Postgresql, 'write_recovery_conf', Mock())
+@patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
 class TestPatroni(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestPatroni, self).__init__(method_name)
-
-    def set_up(self):
+    def setUp(self):
         self.touched = False
         self.init_cancelled = False
-        subprocess.call = subprocess_call
-        psycopg2.connect = psycopg2_connect
-        self.time_sleep = time.sleep
-        time.sleep = nop
-        self.write_pg_hba = Postgresql.write_pg_hba
-        self.write_recovery_conf = Postgresql.write_recovery_conf
-        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__is_shut_down = Mock()
         RestApiServer._BaseServer__shutdown_request = True
         RestApiServer.socket = 0
         with open('postgres0.yml', 'r') as f:
@@ -105,50 +69,38 @@ class TestPatroni(unittest.TestCase):
             with patch.object(Client, 'machines') as mock_machines:
                 mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
                 self.p = Patroni(config)
+                self.p.ha.dcs.client.write = etcd_write
+                self.p.ha.dcs.client.read = etcd_read
 
-    def tear_down(self):
-        time.sleep = self.time_sleep
-        Postgresql.write_pg_hba = self.write_pg_hba
-        Postgresql.write_recovery_conf = self.write_recovery_conf
-
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient())
     def test_get_dcs(self):
-        patroni.zookeeper.KazooClient = MockKazooClient
         self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
         self.assertRaises(Exception, self.p.get_dcs, '', {})
 
+    @patch('time.sleep', Mock(side_effect=SleepException()))
+    @patch.object(Patroni, 'initialize', Mock())
+    @patch.object(Etcd, 'delete_leader', Mock())
     def test_patroni_main(self):
         main()
         sys.argv = ['patroni.py', 'postgres0.yml']
-        time.sleep = time_sleep
 
         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(SleepException, main)
-
-            Patroni.run = keyboard_interrupt
-            main()
-
-            Patroni.run = run
-            Patroni.touch_member = touch_member
+            with patch.object(Patroni, 'touch_member', self.touch_member):
+                with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
+                    self.assertRaises(SleepException, main)
+                with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
+                    main()
 
+    @patch('time.sleep', Mock(side_effect=SleepException()))
     def test_patroni_run(self):
-        time.sleep = time_sleep
         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.watch = time_sleep
         self.assertRaises(SleepException, self.p.run)
+
         self.p.ha.state_handler.is_leader = false
-        self.p.api.start = nop
+        self.p.api.start = Mock()
         self.assertRaises(SleepException, self.p.run)
 
     def touch_member(self, ttl=None):
@@ -158,7 +110,6 @@ class TestPatroni(unittest.TestCase):
         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.p.postgresql.name, 'b', 'c', (now + datetime.timedelta(
@@ -167,8 +118,6 @@ class TestPatroni(unittest.TestCase):
         self.p.touch_member()
 
     def test_patroni_initialize(self):
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.ha.dcs.client.read = etcd_read
         self.p.touch_member = self.touch_member
         self.p.postgresql.data_directory_empty = true
         self.p.ha.dcs.initialize = true
@@ -179,25 +128,24 @@ class TestPatroni(unittest.TestCase):
 
         self.p.ha.dcs.initialize = false
         self.p.ha.dcs.get_cluster = get_cluster_initialized_with_leader
-        time.sleep = time_sleep
-        self.p.ha.dcs.client.read = etcd_read
-        self.p.initialize()
+        with patch('time.sleep', time_sleep):
+            self.p.initialize()
 
-        self.p.ha.dcs.get_cluster = get_cluster_initialized_without_leader
-        self.assertRaises(SleepException, self.p.initialize)
+            self.p.ha.dcs.get_cluster = get_cluster_initialized_without_leader
+            self.assertRaises(SleepException, self.p.initialize)
 
-        self.p.postgresql.data_directory_empty = false
-        self.p.initialize()
+            self.p.postgresql.data_directory_empty = false
+            self.p.initialize()
 
-        self.p.ha.dcs.get_cluster = get_cluster_not_initialized_with_leader
-        self.p.postgresql.data_directory_empty = true
-        self.p.initialize()
+            self.p.ha.dcs.get_cluster = get_cluster_not_initialized_with_leader
+            self.p.postgresql.data_directory_empty = true
+            self.p.initialize()
 
-        self.p.ha.dcs.get_cluster = get_cluster_dcs_error
-        self.assertRaises(SleepException, self.p.initialize)
+            self.p.ha.dcs.get_cluster = get_cluster_dcs_error
+            self.assertRaises(SleepException, self.p.initialize)
 
     def test_schedule_next_run(self):
-        self.p.ha.dcs.watch = lambda e: True
+        self.p.ha.dcs.watch = Mock(return_value=True)
         self.p.schedule_next_run()
         self.p.next_run = time.time() - self.p.nap_time - 1
         self.p.schedule_next_run()
@@ -206,8 +154,6 @@ class TestPatroni(unittest.TestCase):
         self.init_cancelled = True
 
     def test_cleanup_on_initialization(self):
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.ha.dcs.client.read = etcd_read
         self.p.ha.dcs.get_cluster = get_cluster_not_initialized_without_leader
         self.p.touch_member = self.touch_member
         self.p.postgresql.data_directory_empty = true

tests/test_postgresql.py

@@ -1,33 +1,28 @@
 import os
 import psycopg2
 import shutil
-import subprocess
 import unittest
 
+from mock import Mock, patch
 from patroni.dcs import Cluster, Leader, Member
+from patroni.exceptions import PostgresConnectionException
 from patroni.postgresql import Postgresql
-from test_ha import true, false
-
-
-def nop(*args, **kwargs):
-    pass
-
-
-def subprocess_call(cmd, shell=False, env=None):
-    return 0
+from patroni.utils import RetryFailedError
+from test_ha import false
 
 
 class MockCursor:
 
-    def __init__(self):
+    def __init__(self, connection):
+        self.connection = connection
         self.closed = False
         self.results = []
 
     def execute(self, sql, *params):
         if sql.startswith('blabla') or sql == 'CHECKPOINT':
             raise psycopg2.OperationalError()
-        elif sql.startswith('InterfaceError'):
-            raise psycopg2.InterfaceError()
+        elif sql.startswith('RetryFailedError'):
+            raise RetryFailedError('retry')
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
         elif sql.startswith('SELECT pg_current_xlog_location()'):
@@ -69,38 +64,32 @@ class MockCursor:
         for i in self.results:
             yield i
 
+    def __enter__(self):
+        return self
 
-class MockConnect:
+    def __exit__(self, *args):
+        pass
 
-    def __init__(self):
-        self.autocommit = False
-        self.closed = 0
+
+class MockConnect(Mock):
+
+    autocommit = False
+    closed = 0
 
     def cursor(self):
-        return MockCursor()
-
-    def close(self):
-        pass
+        return MockCursor(self)
 
 
 def psycopg2_connect(*args, **kwargs):
     return MockConnect()
 
 
-def raise_exception(*args, **kwargs):
-    raise Exception
-
-
+@patch('subprocess.call', Mock(return_value=0))
+@patch('shutil.copy', Mock())
+@patch('psycopg2.connect', psycopg2_connect)
 class TestPostgresql(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestPostgresql, self).__init__(method_name)
-
-    def set_up(self):
-        subprocess.call = subprocess_call
-        shutil.copy = nop
+    def setUp(self):
         self.p = Postgresql({'name': 'test0', 'scope': 'batman', 'data_dir': 'data/test0',
                              'listen': '127.0.0.1, *:5432', 'connect_address': '127.0.0.2:5432',
                              'pg_hba': ['hostssl all all 0.0.0.0/0 md5', 'host all all 0.0.0.0/0 md5'],
@@ -115,7 +104,6 @@ class TestPostgresql(unittest.TestCase):
                                            'on_reload': 'true'
                                            },
                              'restore': 'true'})
-        psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)
         self.leadermem = Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres', None, None, 28)
@@ -123,12 +111,9 @@ class TestPostgresql(unittest.TestCase):
         self.other = Member(0, 'test1', 'postgres://replicator:[email protected]:5433/postgres', None, None, 28)
         self.me = Member(0, 'test0', 'postgres://replicator:[email protected]:5434/postgres', None, None, 28)
 
-    def tear_down(self):
+    def tearDown(self):
         shutil.rmtree('data')
 
-    def mock_query(self, p):
-        raise psycopg2.OperationalError("not supported")
-
     def test_data_directory_empty(self):
         self.assertTrue(self.p.data_directory_empty())
 
@@ -155,7 +140,7 @@ class TestPostgresql(unittest.TestCase):
         self.p.follow_the_leader(Leader(-1, None, 28, self.other))
 
     def test_create_replica(self):
-        self.p.delete_trigger_file = raise_exception
+        self.p.delete_trigger_file = Mock(side_effect=OSError())
         self.assertEquals(self.p.create_replica({'host': '', 'port': '', 'user': ''}, ''), 1)
 
     def test_create_connection_users(self):
@@ -169,14 +154,13 @@ class TestPostgresql(unittest.TestCase):
         cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leadermem])
         self.p.sync_replication_slots(cluster)
 
+    @patch.object(MockConnect, 'closed', 2)
+    def test__query(self):
+        self.assertRaises(PostgresConnectionException, self.p._query, 'blabla')
+
     def test_query(self):
         self.p.query('select 1')
-        self.assertRaises(psycopg2.InterfaceError, self.p.query, 'InterfaceError')
-        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p._connection.closed = 2
-        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p._connection.closed = 2
-        self.p.disconnect = false
+        self.assertRaises(PostgresConnectionException, self.p.query, 'RetryFailedError')
         self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
 
     def test_is_healthiest_node(self):
@@ -206,24 +190,22 @@ class TestPostgresql(unittest.TestCase):
     def test_last_operation(self):
         self.assertEquals(self.p.last_operation(), '0')
 
+    @patch('subprocess.Popen', Mock(side_effect=OSError()))
     def test_call_nowait(self):
-        popen = subprocess.Popen
-        subprocess.Popen = raise_exception
         self.assertFalse(self.p.call_nowait('on_start'))
-        subprocess.Popen = popen
 
     def test_non_existing_callback(self):
         self.assertFalse(self.p.call_nowait('foobar'))
 
     def test_is_leader_exception(self):
         self.p.start()
-        self.p.query = self.mock_query
+        self.p.query = Mock(side_effect=psycopg2.OperationalError("not supported"))
         self.assertTrue(self.p.stop())
 
+    @patch('os.rename', Mock())
+    @patch('os.path.isdir', Mock(return_value=True))
     def test_move_data_directory(self):
         self.p.is_running = false
-        os.rename = nop
-        os.path.isdir = true
-        self.p.move_data_directory()
-        os.rename = raise_exception
         self.p.move_data_directory()
+        with patch('os.rename', Mock(side_effect=OSError())):
+            self.p.move_data_directory()

tests/test_utils.py

@@ -1,56 +1,34 @@
-import os
-import time
 import unittest
 
-from patroni.exceptions import DCSError
+from mock import Mock, patch
+from patroni.exceptions import PatroniException
 from patroni.utils import Retry, RetryFailedError, reap_children, sigchld_handler, sigterm_handler, sleep
 
 
-def nop(*args, **kwargs):
-    pass
-
-
-def os_waitpid(a, b):
-    return (0, 0)
-
-
 def time_sleep(_):
     sigchld_handler(None, None)
 
 
 class TestUtils(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestUtils, self).__init__(method_name)
-
-    def set_up(self):
-        self.time_sleep = time.sleep
-        time.sleep = nop
-
-    def tear_down(self):
-        time.sleep = self.time_sleep
-
     def test_sigterm_handler(self):
         self.assertRaises(SystemExit, sigterm_handler, None, None)
 
+    @patch('time.sleep', Mock())
     def test_reap_children(self):
         reap_children()
-        os.waitpid = os_waitpid
-        sigchld_handler(None, None)
-        reap_children()
+        with patch('os.waitpid', Mock(return_value=(0, 0))):
+            sigchld_handler(None, None)
+            reap_children()
 
+    @patch('time.sleep', time_sleep)
     def test_sleep(self):
-        time.sleep = time_sleep
         sleep(0.01)
 
 
+@patch('time.sleep', Mock())
 class TestRetrySleeper(unittest.TestCase):
 
-    def _pass(self):
-        pass
-
     def _fail(self, times=1):
         scope = dict(times=0)
 
@@ -59,7 +37,7 @@ class TestRetrySleeper(unittest.TestCase):
                 pass
             else:
                 scope['times'] += 1
-                raise DCSError('Failed!')
+                raise PatroniException('Failed!')
         return inner
 
     def _makeOne(self, *args, **kwargs):
@@ -78,20 +56,14 @@ class TestRetrySleeper(unittest.TestCase):
         self.assertEquals(retry._attempts, 1)
 
     def test_maximum_delay(self):
-        def sleep_func(_time):
-            pass
-
-        retry = self._makeOne(delay=10, max_tries=100, sleep_func=sleep_func)
+        retry = self._makeOne(delay=10, max_tries=100)
         retry(self._fail(times=10))
         self.assertTrue(retry._cur_delay < 4000, retry._cur_delay)
         # gevent's sleep function is picky about the type
         self.assertEquals(type(retry._cur_delay), float)
 
     def test_deadline(self):
-        def sleep_func(_time):
-            pass
-
-        retry = self._makeOne(deadline=0.0001, sleep_func=sleep_func)
+        retry = self._makeOne(deadline=0.0001)
         self.assertRaises(RetryFailedError, retry, self._fail(times=100))
 
     def test_copy(self):

tests/test_zookeeper.py

@@ -1,58 +1,25 @@
-import patroni.zookeeper
-import requests
 import six
 import unittest
 
+from mock import Mock, patch
 from patroni.dcs import Leader
 from patroni.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError
 from kazoo.client import KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
 from kazoo.protocol.states import ZnodeStat
-from test_etcd import MockPostgresql, requests_get
+from test_etcd import MockPostgresql, SleepException, requests_get
 
 
-class MockEvent:
+class MockKazooClient(Mock):
 
-    def clear(self):
-        pass
-
-    def set(self):
-        pass
-
-    def wait(self, timeout):
-        pass
-
-    def isSet(self):
-        return True
-
-
-class MockEventHandler:
-
-    def event_object(self):
-        return MockEvent()
-
-
-class SleepException(Exception):
-    pass
-
-
-class MockKazooClient:
-
-    def __init__(self, **kwargs):
-        self.handler = MockEventHandler()
-        self.leader = False
-        self.exists = True
-
-    def start(self, timeout):
-        pass
+    leader = False
+    exists = True
+    handler = Mock()
 
     @property
     def client_id(self):
         return (-1, '')
 
-    def add_listener(self, cb):
-        pass
-
     def retry(self, func, *args, **kwargs):
         func(*args, **kwargs)
 
@@ -115,37 +82,20 @@ class MockKazooClient:
         elif path.endswith('/initialize'):
             raise NoNodeError
 
-    def set_hosts(self, hosts, randomize_hosts=None):
-        pass
-
-
-def exhibitor_sleep(_):
-    raise SleepException
-
 
+@patch('requests.get', requests_get)
+@patch('patroni.zookeeper.sleep', Mock(side_effect=SleepException()))
 class TestExhibitorEnsembleProvider(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestExhibitorEnsembleProvider, self).__init__(method_name)
-
-    def set_up(self):
-        requests.get = requests_get
-        patroni.zookeeper.sleep = exhibitor_sleep
-
     def test_init(self):
         self.assertRaises(SleepException, ExhibitorEnsembleProvider, ['localhost'], 8181)
 
 
 class TestZooKeeper(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestZooKeeper, self).__init__(method_name)
-
-    def set_up(self):
-        requests.get = requests_get
-        patroni.zookeeper.KazooClient = MockKazooClient
+    @patch('requests.get', requests_get)
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient)
+    def setUp(self):
         self.zk = ZooKeeper('foo', {'exhibitor': {'hosts': ['localhost', 'exhibitor'], 'port': 8181}, 'scope': 'test'})
 
     def test_session_listener(self):