From 9f9cb6005d7baed8b18e897fb23c84281980a680 Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Sun, 20 Sep 2015 13:41:57 +0200 Subject: [PATCH 1/2] Refactor Postgresql.query method to use commont retry mechanism query method in an api.py also needs retry in some cases (for example when we are running is_healthiest_node check). In all cases We will retry only when connection is closed or broken, BUT connection status must be checked via cursor.connection (old implementation was using general connection object for that). For multi-threaded applications this is not appropriate, because some other thread might restore connection. In appdition to that I've changed most of the unit tests to use `Mock` and `patch` where it is possible. --- patroni/api.py | 50 +++++++++++----- patroni/exceptions.py | 6 +- patroni/ha.py | 4 +- patroni/postgresql.py | 41 ++++++------- patroni/utils.py | 15 ++--- tests/test_api.py | 63 +++++++------------- tests/test_etcd.py | 59 ++++++------------ tests/test_ha.py | 23 ++----- tests/test_patroni.py | 126 +++++++++++---------------------------- tests/test_postgresql.py | 88 +++++++++++---------------- tests/test_utils.py | 50 ++++------------ tests/test_zookeeper.py | 72 ++++------------------ 12 files changed, 205 insertions(+), 392 deletions(-) diff --git a/patroni/api.py b/patroni/api.py index e249c5d4..30bc8914 100644 --- a/patroni/api.py +++ b/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): diff --git a/patroni/exceptions.py b/patroni/exceptions.py index 507edfc3..97985696 100644 --- a/patroni/exceptions.py +++ b/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 diff --git a/patroni/ha.py b/patroni/ha.py index 7725d3c2..9ed8faf3 100644 --- a/patroni/ha.py +++ b/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') diff --git a/patroni/postgresql.py b/patroni/postgresql.py index 398a8ab8..532e65c8 100644 --- a/patroni/postgresql.py +++ b/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) == [] diff --git a/patroni/utils.py b/patroni/utils.py index 19d5837a..681ca39c 100644 --- a/patroni/utils.py +++ b/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) diff --git a/tests/test_api.py b/tests/test_api.py index d70d87a4..b8f9e038 100644 --- a/tests/test_api.py +++ b/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') diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 9ded883b..d2824699 100644 --- a/tests/test_etcd.py +++ b/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): diff --git a/tests/test_ha.py b/tests/test_ha.py index d82668e6..3fa50b41 100644 --- a/tests/test_ha.py +++ b/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') diff --git a/tests/test_patroni.py b/tests/test_patroni.py index 7481a6fd..13621bf9 100644 --- a/tests/test_patroni.py +++ b/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 diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 7e8e768f..eaf84943 100644 --- a/tests/test_postgresql.py +++ b/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:rep-pass@127.0.0.1:5435/postgres', None, None, 28) @@ -123,12 +111,9 @@ class TestPostgresql(unittest.TestCase): self.other = Member(0, 'test1', 'postgres://replicator:rep-pass@127.0.0.1:5433/postgres', None, None, 28) self.me = Member(0, 'test0', 'postgres://replicator:rep-pass@127.0.0.1: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() diff --git a/tests/test_utils.py b/tests/test_utils.py index 9a814f78..45b194dc 100644 --- a/tests/test_utils.py +++ b/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): diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py index f2141707..afbdad5f 100644 --- a/tests/test_zookeeper.py +++ b/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): From d8982e1e5a811d06690433a9c4afd86664f681ef Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Sun, 20 Sep 2015 13:41:57 +0200 Subject: [PATCH 2/2] Refactor Postgresql.query method to use common retry mechanism query method in an api.py also needs retry in some cases (for example when we are running is_healthiest_node check). In all cases we should retry only when connection is closed or broken. BUT, the connection status must be checked via cursor.connection (old implementation was using general connection object for that). For multi-threaded applications this is not appropriate, because some other thread might restore connection. In addition to that I've changed most of the unit tests to use `Mock` and `patch` where it is possible. --- patroni/api.py | 50 +++++++++++----- patroni/exceptions.py | 6 +- patroni/ha.py | 4 +- patroni/postgresql.py | 41 ++++++------- patroni/utils.py | 15 ++--- tests/test_api.py | 63 +++++++------------- tests/test_etcd.py | 59 ++++++------------ tests/test_ha.py | 23 ++----- tests/test_patroni.py | 126 +++++++++++---------------------------- tests/test_postgresql.py | 88 +++++++++++---------------- tests/test_utils.py | 50 ++++------------ tests/test_zookeeper.py | 72 ++++------------------ 12 files changed, 205 insertions(+), 392 deletions(-) diff --git a/patroni/api.py b/patroni/api.py index e249c5d4..30bc8914 100644 --- a/patroni/api.py +++ b/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): diff --git a/patroni/exceptions.py b/patroni/exceptions.py index 507edfc3..97985696 100644 --- a/patroni/exceptions.py +++ b/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 diff --git a/patroni/ha.py b/patroni/ha.py index 7725d3c2..9ed8faf3 100644 --- a/patroni/ha.py +++ b/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') diff --git a/patroni/postgresql.py b/patroni/postgresql.py index 398a8ab8..532e65c8 100644 --- a/patroni/postgresql.py +++ b/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) == [] diff --git a/patroni/utils.py b/patroni/utils.py index 19d5837a..681ca39c 100644 --- a/patroni/utils.py +++ b/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) diff --git a/tests/test_api.py b/tests/test_api.py index d70d87a4..b8f9e038 100644 --- a/tests/test_api.py +++ b/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') diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 9ded883b..d2824699 100644 --- a/tests/test_etcd.py +++ b/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): diff --git a/tests/test_ha.py b/tests/test_ha.py index d82668e6..3fa50b41 100644 --- a/tests/test_ha.py +++ b/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') diff --git a/tests/test_patroni.py b/tests/test_patroni.py index 7481a6fd..13621bf9 100644 --- a/tests/test_patroni.py +++ b/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 diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 7e8e768f..eaf84943 100644 --- a/tests/test_postgresql.py +++ b/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:rep-pass@127.0.0.1:5435/postgres', None, None, 28) @@ -123,12 +111,9 @@ class TestPostgresql(unittest.TestCase): self.other = Member(0, 'test1', 'postgres://replicator:rep-pass@127.0.0.1:5433/postgres', None, None, 28) self.me = Member(0, 'test0', 'postgres://replicator:rep-pass@127.0.0.1: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() diff --git a/tests/test_utils.py b/tests/test_utils.py index 9a814f78..45b194dc 100644 --- a/tests/test_utils.py +++ b/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): diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py index f2141707..afbdad5f 100644 --- a/tests/test_zookeeper.py +++ b/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):