import psycopg2 import unittest from mock import Mock, patch from patroni.api import RestApiHandler, RestApiServer from patroni.dcs import Member from six import BytesIO as IO from six.moves import BaseHTTPServer from six.moves.BaseHTTPServer import BaseHTTPRequestHandler import socket from test_postgresql import psycopg2_connect, MockCursor class MockPostgresql(Mock): name = 'test' state = 'running' role = 'master' server_version = '999999' scope = 'dummy' @staticmethod def connection(): return psycopg2_connect() @staticmethod def is_running(): return True class MockHa(Mock): dcs = Mock() state_handler = MockPostgresql() @staticmethod def schedule_restart(): return 'restart' @staticmethod def schedule_reinitialize(): return 'reinitialize' @staticmethod def restart(): return (True, '') @staticmethod def restart_scheduled(): return False @staticmethod def fetch_nodes_statuses(members): return [[None, True, None, None, {}]] class MockPatroni(Mock): postgresql = MockPostgresql() ha = MockHa() dcs = Mock() tags = {} version = '0.00' class MockRequest(object): def __init__(self, path): self.path = path def makefile(self, *args, **kwargs): return IO(self.path) class MockRestApiServer(RestApiServer): 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) Handler(MockRequest(path), ('0.0.0.0', 8080), self) @patch('ssl.wrap_socket', Mock(return_value=0)) class TestRestApiHandler(unittest.TestCase): def test_do_GET(self): MockRestApiServer(RestApiHandler, b'GET /replica') with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={})): MockRestApiServer(RestApiHandler, b'GET /replica') with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})): MockRestApiServer(RestApiHandler, b'GET /replica') MockRestApiServer(RestApiHandler, b'GET /master') MockPatroni.dcs.cluster.leader.name = MockPostgresql.name MockRestApiServer(RestApiHandler, b'GET /replica') MockPatroni.dcs.cluster = None with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})): MockRestApiServer(RestApiHandler, b'GET /master') with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)): MockRestApiServer(RestApiHandler, b'GET /master') MockRestApiServer(RestApiHandler, b'GET /master') def test_do_OPTIONS(self): MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0') with patch.object(BaseHTTPRequestHandler, 'handle_one_request') as mock_handle_request: mock_handle_request.side_effect = socket.error("foo") MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0') # make sure socket.error gets propagated via wfile object in finalize() with patch.object(MockRequest, 'makefile') as makefile: makefile.return_value.closed = False makefile.return_value.readline.side_effect = lambda x: b"foo" makefile.return_value.flush = Mock(side_effect=socket.error("foo")) MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0') def test_do_GET_patroni(self): MockRestApiServer(RestApiHandler, b'GET /patroni') def test_basicauth(self): MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0') MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0\nAuthorization:') def test_do_POST_restart(self): request = b'POST /restart HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0' MockRestApiServer(RestApiHandler, request) with patch.object(MockHa, 'restart', Mock(side_effect=Exception)): MockRestApiServer(RestApiHandler, request) @patch.object(MockHa, 'dcs') def test_do_POST_reinitialize(self, dcs): cluster = dcs.get_cluster.return_value request = b'POST /reinitialize HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0' MockRestApiServer(RestApiHandler, request) cluster.is_unlocked.return_value = False MockRestApiServer(RestApiHandler, request) with patch.object(MockHa, 'schedule_reinitialize', Mock(return_value=None)): MockRestApiServer(RestApiHandler, request) cluster.leader.name = 'test' MockRestApiServer(RestApiHandler, request) @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') @patch('time.sleep', Mock()) @patch.object(MockHa, 'dcs') def test_do_POST_failover(self, dcs): cluster = dcs.get_cluster.return_value request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\ b'Content-Length: 25\n\n{"leader": "postgresql1"}' MockRestApiServer(RestApiHandler, request) cluster.leader.name = 'postgresql1' MockRestApiServer(RestApiHandler, request) cluster.members = [Member(0, 'postgresql0', 30, {'api_url': 'http'})] MockRestApiServer(RestApiHandler, request) with patch.object(MockPatroni, 'dcs') as d: cluster = d.get_cluster.return_value cluster.leader.name = 'postgresql0' MockRestApiServer(RestApiHandler, request) cluster.leader.name = 'postgresql1' cluster.failover = None MockRestApiServer(RestApiHandler, request) d.get_cluster = Mock(side_effect=Exception()) MockRestApiServer(RestApiHandler, request) d.manual_failover.return_value = False MockRestApiServer(RestApiHandler, request) with patch.object(MockHa, 'fetch_nodes_statuses', Mock(return_value=[])): MockRestApiServer(RestApiHandler, request) request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\ b'Content-Length: 50\n\n{"leader": "postgresql1", "member": "postgresql2"}' MockRestApiServer(RestApiHandler, request)