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):