From befd33555d168fc78b2fe3fc040771a942cebe66 Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Mon, 24 Aug 2015 16:58:08 +0200 Subject: [PATCH] Refactor helpers/etcd.py Work with etcd cluster via high-level python-etcd module. Plus change all unit tests accordingly. --- helpers/etcd.py | 348 ++++++++++++++------------------------- patroni.py | 8 +- requirements-py2.txt | 1 + requirements-py3.txt | 1 + tests/test_api.py | 3 +- tests/test_etcd.py | 194 +++++++++++----------- tests/test_ha.py | 25 +-- tests/test_patroni.py | 102 ++++++++---- tests/test_postgresql.py | 3 - tests/test_zookeeper.py | 2 +- 10 files changed, 314 insertions(+), 373 deletions(-) diff --git a/helpers/etcd.py b/helpers/etcd.py index 83f9e0a9..dc9f410d 100644 --- a/helpers/etcd.py +++ b/helpers/etcd.py @@ -1,4 +1,7 @@ +from __future__ import absolute_import +import etcd import logging +import os import random import requests import socket @@ -16,49 +19,37 @@ class EtcdError(DCSError): pass -class EtcdConnectionFailed(EtcdError): - pass - - -class Client: - - API_VERSION = 'v2' +class Client(etcd.Client): def __init__(self, config): + super(Client, self).__init__(read_timeout=5) self._config = config - self.timeout = 5 - self._base_uri = None - self._members_cache = [] - self.load_members() + self._load_machines_cache() + self._allow_reconnect = True - def client_url(self, path): - return self._base_uri + path - - def _next_server(self): - self._base_uri = None + @property + def machines(self): + """Original `machines` method(property) of `etcd.Client` class raise exception + when it failed to get list of etcd cluster members. This method is being called + only when request failed on one of the etcd members during `api_execute` call. + For us it's more important to execute original request rather then get new + topology of etcd cluster. So we will catch this exception and return valid list + of machines with setting flag `self._update_machines_cache` to `!True`. + Later, during next `api_execute` call we will forcefully update machines_cache""" try: - self._base_uri = self._members_cache.pop() - except IndexError: - logger.error('Members cache is empty, can not retry.') - raise EtcdConnectionFailed('No more members in the cluster') - else: - logger.info('Selected new etcd server %s', self._base_uri) + ret = super(Client, self).machines + random.shuffle(ret) + return ret + except etcd.EtcdException: + if self._update_machines_cache: # We are updating machines_cache + raise # This exception is fatal, we should re-raise it. + self._update_machines_cache = True + return [self._base_uri] - def _get(self, path): - response = None - while response is None: - uri = self.client_url(path) - try: - logger.info('GET %s', uri) - response = requests.get(uri, timeout=self.timeout) - except RequestException: - self._next_server() - - logger.debug([response.status_code, response.content]) - try: - return response.json(), response.status_code - except (TypeError, ValueError): - raise EtcdError('Bad response from %s: %s' % (uri, response.content)) + def api_execute(self, path, method, **kwargs): + # Update machines_cache if previous attempt of update has failed + self._update_machines_cache and self._load_machines_cache() + return super(Client, self).api_execute(path, method, **kwargs) @staticmethod def get_srv_record(host): @@ -68,120 +59,72 @@ class Client: logger.exception('Can not resolve SRV for %s', host) return [] - @staticmethod - def get_peers_urls_from_dns(host): - return ['http://{}:{}'.format(h, p) for h, p in Client.get_srv_record(host)] + def _get_machines_cache_from_srv(self, discovery_srv): + """Fetch list of etcd-cluster member by resolving _etcd-server._tcp. SRV record. + This record should contain list of host and peer ports which could be used to run + 'GET http://{host}:{port}/members' request (peer protocol)""" - @staticmethod - def get_client_urls_from_dns(addr): - host, port = addr.split(':') ret = [] + for host, port in self.get_srv_record(discovery_srv): + url = '{}://{}:{}/members'.format(self._protocol, host, port) + try: + response = requests.get(url) + if response.ok: + for member in response.json(): + ret.extend(member['clientURLs']) + break + except RequestException: + logger.exception('GET %s', url) + return list(set(ret)) + + def _get_machines_cache_from_dns(self, addr): + """One host might be resolved into multiple ip addresses. We will make list out of it""" + + ret = [] + host, port = addr.split(':') try: for r in set(socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)): - ret.append('http://{}:{}/{}'.format(r[4][0], r[4][1], Client.API_VERSION)) + ret.append('{}://{}:{}'.format(self._protocol, r[4][0], r[4][1])) except socket.error: logger.exception('Can not resolve %s', host) - return list(set(ret)) if ret else ['http://{}:{}/{}'.format(host, port, Client.API_VERSION)] + return list(set(ret)) if ret else ['{}://{}:{}'.format(self._protocol, host, port)] - def load_members(self): - load_from_srv = False - if not self._base_uri: - if 'discovery_srv' not in self._config and 'host' not in self._config: - raise Exception('Neither discovery_srv nor host are defined in etcd section of config') + def _load_machines_cache(self): + """This method should fill up `_machines_cache` from scratch. + It could happen only in two cases: + 1. During class initialization + 2. When all etcd members failed""" - if 'discovery_srv' in self._config: - load_from_srv = True - self._members_cache = self.get_peers_urls_from_dns(self._config['discovery_srv']) + self._update_machines_cache = True - if not self._members_cache and 'host' in self._config: - load_from_srv = False - self._members_cache = self.get_client_urls_from_dns(self._config['host']) + if 'discovery_srv' not in self._config and 'host' not in self._config: + raise Exception('Neither discovery_srv nor host are defined in etcd section of config') - self._next_server() + if 'discovery_srv' in self._config: + self._machines_cache = self._get_machines_cache_from_srv(self._config['discovery_srv']) - response, status_code = self._get('/members') - if status_code != 200: - self._base_uri = None - raise EtcdError('Got response with code=%s from %s' % (status_code, self._base_uri)) + if not self._machines_cache and 'host' in self._config: + self._machines_cache = self._get_machines_cache_from_dns(self._config['host']) - members_cache = [] + # Can not bootstrap list of etcd-cluster members, giving up + if not self._machines_cache: + raise etcd.EtcdException + + # After filling up initial list of machines_cache we should ask etcd-cluster about actual list + self._base_uri = self._machines_cache.pop(0) + self._machines_cache = self.machines + self._base_uri in self._machines_cache and self._machines_cache.remove(self._base_uri) + + self._update_machines_cache = False + + +def catch_etcd_errors(func): + def wrapper(*args, **kwargs): try: - for member in response if load_from_srv else response['members']: - members_cache.extend([m + '/' + self.API_VERSION for m in member['clientURLs']]) - except: - self._base_uri = None - raise EtcdError('Got invalid response from %s: %s' % (self._base_uri, response)) - - self._members_cache = list(set(members_cache)) - random.shuffle(self._members_cache) - if load_from_srv: - self._next_server() - else: - try: - self._members_cache.remove(self._base_uri) - except ValueError: - pass - - def get(self, path): - if not self._base_uri: - self.load_members() - old_base_uri = self._base_uri - try: - return self._get(path) - finally: - if self._base_uri != old_base_uri: - try: - self.load_members() - except EtcdError: - logger.exception('load_members') - - def put(self, path, **data): - if not self._base_uri: - self.load_members() - old_base_uri = self._base_uri - response = None - while response is None: - uri = self.client_url(path) - try: - logger.info('PUT %s', uri) - response = requests.put(uri, timeout=self.timeout, data=data) - except RequestException: - logger.exception('PUT %s data=%s', uri, data) - self._next_server() - - if self._base_uri != old_base_uri: - try: - self.load_members() - except EtcdError: - logger.exception('load_members') - if response.status_code in [200, 201, 202, 204]: - return True - logger.error('Unexpected response: %s %s', response.status_code, response.content) - return False - - def delete(self, path): - if not self._base_uri: - self.load_members() - old_base_uri = self._base_uri - response = None - while response is None: - uri = self.client_url(path) - try: - logger.info('DELETE %s', uri) - response = requests.delete(uri, timeout=self.timeout) - except RequestException: - logger.exception('DELETE %s', uri) - self._next_server() - - if self._base_uri != old_base_uri: - try: - self.load_members() - except EtcdError: - logger.exception('load_members') - if response.status_code in [200, 202, 204]: - return True - logger.error('Unexpected response: %s %s', response.status_code, response.content) - return False + return not func(*args, **kwargs) is None + except etcd.EtcdException: + return False + return wrapper class Etcd(AbstractDCS): @@ -197,118 +140,73 @@ class Etcd(AbstractDCS): while not client: try: client = Client(config) - except EtcdError: + except etcd.EtcdException: logger.info('waiting on etcd') sleep(5) return client - def client_path(self, path): - return '/keys' + super(Etcd, self).client_path(path) - - def get_client_path(self, path): - return self.client.get(self.client_path(path)) - - def put_client_path(self, path, **data): - return self.client.put(self.client_path(path), **data) - - def delete_client_path(self, path): - try: - return self.client.delete(self.client_path(path)) - except EtcdConnectionFailed: - return False - - @staticmethod - def find_node(node, key): - """ - >>> Etcd.find_node({}, None) - >>> Etcd.find_node({'dir': True, 'nodes': [], 'key': '/test/'}, 'test') - """ - if not node.get('dir', False): - return None - key = node['key'] + key - for n in node['nodes']: - if n['key'] == key: - return n - return None - @staticmethod def member(node): - conn_url, api_url = parse_connection_string(node['value']) - expiration = node.get('expiration', None) - ttl = node.get('ttl', None) - return Member(node['modifiedIndex'], node['key'].split('/')[-1], conn_url, api_url, expiration, ttl) + conn_url, api_url = parse_connection_string(node.value) + return Member(node.modifiedIndex, os.path.basename(node.key), conn_url, api_url, node.expiration, node.ttl) def get_cluster(self): try: - response, status_code = self.get_client_path('?recursive=true') - if status_code == 200: - node = self.find_node(response['node'], '/initialize') - initialize = True if node else False - # get list of members - node = self.find_node(response['node'], '/members') or {'nodes': []} - members = [self.member(n) for n in node['nodes']] + result = self.client.read(self.client_path(''), recursive=True) + nodes = {os.path.relpath(node.key, result.key): node for node in result.leaves} - # get last leader operation - last_leader_operation = 0 - node = self.find_node(response['node'], '/optime') - if node: - node = self.find_node(node, '/leader') - if node: - last_leader_operation = int(node['value']) + # get initialize flag + initialize = bool(nodes.get('initialize', False)) - # get leader - leader = None - node = self.find_node(response['node'], '/leader') - if node: - for m in members: - if m.name == node['value']: - leader = m - break - if not leader: - leader = Member(-1, node['value'], None, None, None, None) + # get last leader operation + last_leader_operation = nodes.get('optime/leader', None) + last_leader_operation = 0 if last_leader_operation is None else int(last_leader_operation.value) - return Cluster(initialize, leader, last_leader_operation, members) - elif status_code == 404: - return Cluster(False, None, None, []) + # get list of members + members = [self.member(n) for k, n in nodes.items() if k.startswith('members/') and len(k.split('/')) == 2] + + # get leader + leader = nodes.get('leader', None) + if leader: + leader = Member(-1, leader.value, None, None, None, None) + leader = ([m for m in members if m.name == leader.name] or [leader])[0] + + return Cluster(initialize, leader, last_leader_operation, members) + except etcd.EtcdKeyNotFound: + return Cluster(False, None, None, []) except: logger.exception('get_cluster') raise EtcdError('Etcd is not responding properly') + @catch_etcd_errors def touch_member(self, connection_string, ttl=None): - try: - return self.put_client_path('/members/' + self._name, value=connection_string, ttl=ttl or self.member_ttl) - except EtcdError: - return False + return self.client.set(self.client_path('/members/' + self._name), connection_string, ttl or self.member_ttl) + @catch_etcd_errors def take_leader(self): - try: - return self.put_client_path('/leader', value=self._name, ttl=self.ttl) - except EtcdError: - return False + return self.client.set(self.client_path('/leader'), self._name, self.ttl) + @catch_etcd_errors def attempt_to_acquire_leader(self): - try: - ret = self.put_client_path('/leader', value=self._name, ttl=self.ttl, prevExist=False) - ret or logger.info('Could not take out TTL lock') - return ret - except EtcdError: - return False + ret = self.client.write(self.client_path('/leader'), self._name, ttl=self.ttl, prevExist=False) + ret or logger.info('Could not take out TTL lock') + return ret + @catch_etcd_errors + def write_leader_optime(self, state_handler): + return self.client.set(self.client_path('/optime/leader'), state_handler.last_operation()) + + @catch_etcd_errors def update_leader(self, state_handler): - if self.put_client_path('/leader', value=state_handler.name, ttl=self.ttl, prevValue=state_handler.name): - try: - self.put_client_path('/optime/leader', value=state_handler.last_operation()) - except EtcdError: - pass - return True - return False + ret = self.client.test_and_set(self.client_path('/leader'), self._name, self._name, self.ttl) + ret and self.write_leader_optime(state_handler) + return ret + @catch_etcd_errors def race(self, path): - try: - return self.put_client_path(path, value=self._name, prevExist=False) - except EtcdError: - return False + return self.client.write(self.client_path(path), self._name, prevExist=False) + @catch_etcd_errors def delete_leader(self): - return self.delete_client_path('/leader?prevValue=' + self._name) + return self.client.delete(self.client_path('/leader'), prevValue=self._name) diff --git a/patroni.py b/patroni.py index 6f06a056..e26a4059 100755 --- a/patroni.py +++ b/patroni.py @@ -12,6 +12,8 @@ from helpers.postgresql import Postgresql from helpers.utils import setup_signal_handlers, sleep from helpers.zookeeper import ZooKeeper +logger = logging.getLogger(__name__) + class Patroni: @@ -44,7 +46,7 @@ class Patroni: def initialize(self): # wait for etcd to be available while not self.touch_member(): - logging.info('waiting on DCS') + logger.info('waiting on DCS') sleep(5) # is data directory empty? @@ -82,14 +84,14 @@ class Patroni: while True: self.touch_member() - logging.info(self.ha.run_cycle()) + logger.info(self.ha.run_cycle()) try: if self.ha.state_handler.is_leader(): self.ha.cluster and self.ha.state_handler.create_replication_slots(self.ha.cluster) else: self.ha.state_handler.drop_replication_slots() except: - logging.exception('Exception when changing replication slots') + logger.exception('Exception when changing replication slots') self.schedule_next_run() diff --git a/requirements-py2.txt b/requirements-py2.txt index cb6203c5..d57df720 100644 --- a/requirements-py2.txt +++ b/requirements-py2.txt @@ -6,3 +6,4 @@ PyYAML requests six >= 1.7 kazoo>=2.2.1 +python-etcd>=0.4.1 diff --git a/requirements-py3.txt b/requirements-py3.txt index 1d5a2316..0fd9dfb3 100644 --- a/requirements-py3.txt +++ b/requirements-py3.txt @@ -6,3 +6,4 @@ PyYAML requests six kazoo>=2.2.1 +python-etcd>=0.4.1 diff --git a/tests/test_api.py b/tests/test_api.py index 07ad410b..91b36943 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,10 +1,9 @@ import psycopg2 -import sys import unittest from helpers.api import RestApiHandler, RestApiServer -from test_postgresql import psycopg2_connect from six import BytesIO as IO +from test_postgresql import psycopg2_connect def throws(*args, **kwargs): diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 599b4bb2..38692c52 100644 --- a/tests/test_etcd.py +++ b/tests/test_etcd.py @@ -1,5 +1,6 @@ import datetime import dns.resolver +import etcd import json import requests import socket @@ -7,8 +8,9 @@ import time import unittest from dns.exception import DNSException -from helpers.dcs import Cluster, Member -from helpers.etcd import Client, Etcd, EtcdConnectionFailed, EtcdError +from helpers.dcs import Cluster, DCSError, Member +from helpers.etcd import Client, Etcd +from mock import Mock, patch class MockResponse: @@ -21,6 +23,17 @@ class MockResponse: def json(self): return json.loads(self.content) + @property + def data(self): + return self.content + + @property + def status(self): + return self.status_code + + def getheader(*args): + return '' + class MockPostgresql: name = '' @@ -30,51 +43,64 @@ class MockPostgresql: def requests_get(url, **kwargs): - members = '[{"id":14855829450254237642,"peerURLs":["http://localhost:2380","http://localhost:7001"],"name":"default","clientURLs":["http://localhost:2379","http://localhost:4001"]}]' + members = '[{"id":14855829450254237642,"peerURLs":["http://localhost:2380","http://localhost:7001"],' +\ + '"name":"default","clientURLs":["http://localhost:2379","http://localhost:4001"]}]' response = MockResponse() - if url.endswith('/v2/members'): - response.content = '{"members": ' + members + '}' - if url.startswith('http://error'): - response.status_code = 404 + if url.startswith('http://local'): + raise requests.exceptions.RequestException() elif url.endswith('/members'): if url.startswith('http://error'): response.content = '[{}]' else: response.content = members - elif url.endswith('/bad_response'): - response.content = '{' elif url.startswith('http://exhibitor'): response.content = '{"servers":["127.0.0.1","127.0.0.2","127.0.0.3"],"port":2181}' - elif url.startswith('http://local'): - raise requests.exceptions.RequestException() - elif url.startswith('http://remote') or url.startswith('http://127.0.0.1') or url.startswith('http://error'): - response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql1","value":"postgres://replicator:rep-pass@127.0.0.1:5434/postgres?application_name=http://127.0.0.1:8009/patroni","expiration":"2015-05-15T09:10:59.949384522Z","ttl":21,"modifiedIndex":20727,"createdIndex":20727},{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:rep-pass@127.0.0.1:5433/postgres?application_name=http://127.0.0.1:8008/patroni","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}' - elif url.startswith('http://other'): - response.status_code = 404 - elif url.startswith('http://noleader'): - response.content = '{"action":"get","node":{"key":"/service/batman5","dir":true,"nodes":[{"key":"/service/batman5/initialize","value":"postgresql0","modifiedIndex":1582,"createdIndex":1582},{"key":"/service/batman5/leader","value":"postgresql1","expiration":"2015-05-15T09:11:00.037397538Z","ttl":21,"modifiedIndex":20728,"createdIndex":20434},{"key":"/service/batman5/optime","dir":true,"nodes":[{"key":"/service/batman5/optime/leader","value":"2164261704","modifiedIndex":20729,"createdIndex":20729}],"modifiedIndex":20437,"createdIndex":20437},{"key":"/service/batman5/members","dir":true,"nodes":[{"key":"/service/batman5/members/postgresql0","value":"postgres://replicator:rep-pass@127.0.0.1:5433/postgres?application_name=http://127.0.0.1:8008/patroni","expiration":"2015-05-15T09:11:09.611860899Z","ttl":30,"modifiedIndex":20730,"createdIndex":20730}],"modifiedIndex":1581,"createdIndex":1581}],"modifiedIndex":1581,"createdIndex":1581}}' else: response.status_code = 404 response.ok = False return response -def requests_put(url, **kwargs): - if url.startswith('http://local') or '/optime/leader' in url: - raise requests.exceptions.RequestException() - response = MockResponse() - response.status_code = 201 - if url.startswith('http://other'): - response.status_code = 404 - return response +def etcd_write(key, value, **kwargs): + if key == '/service/test/leader': + if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True): + return True + raise etcd.EtcdException -def requests_delete(url, **kwargs): - if url.startswith('http://local'): - raise requests.exceptions.RequestException() - response = MockResponse() - response.status_code = 503 if url.startswith('http://error') else 204 - return response +def etcd_delete(key, **kwargs): + raise etcd.EtcdException + + +def etcd_read(key, **kwargs): + if key == '/service/noleader': + raise DCSError('noleader') + elif key == '/service/nocluster': + raise etcd.EtcdKeyNotFound + + response = {"action": "get", "node": {"key": "/service/batman5", "dir": True, "nodes": [ + {"key": "/service/batman5/initialize", "value": "postgresql0", + "modifiedIndex": 1582, "createdIndex": 1582}, + {"key": "/service/batman5/leader", "value": "postgresql1", + "expiration": "2015-05-15T09:11:00.037397538Z", "ttl": 21, + "modifiedIndex": 20728, "createdIndex": 20434}, + {"key": "/service/batman5/optime", "dir": True, "nodes": [ + {"key": "/service/batman5/optime/leader", "value": "2164261704", + "modifiedIndex": 20729, "createdIndex": 20729}], + "modifiedIndex": 20437, "createdIndex": 20437}, + {"key": "/service/batman5/members", "dir": True, "nodes": [ + {"key": "/service/batman5/members/postgresql1", + "value": "postgres://replicator:rep-pass@127.0.0.1:5434/postgres" + + "?application_name=http://127.0.0.1:8009/patroni", + "expiration": "2015-05-15T09:10:59.949384522Z", "ttl": 21, + "modifiedIndex": 20727, "createdIndex": 20727}, + {"key": "/service/batman5/members/postgresql0", + "value": "postgres://replicator:rep-pass@127.0.0.1:5433/postgres" + + "?application_name=http://127.0.0.1:8008/patroni", + "expiration": "2015-05-15T09:11:09.611860899Z", "ttl": 30, + "modifiedIndex": 20730, "createdIndex": 20730}], + "modifiedIndex": 1581, "createdIndex": 1581}], "modifiedIndex": 1581, "createdIndex": 1581}} + return etcd.EtcdResult(**response) def time_sleep(_): @@ -104,6 +130,12 @@ def socket_getaddrinfo(*args): raise socket.error() +def http_request(method, url, **kwargs): + if url == 'http://localhost:2379/': + return MockResponse() + raise socket.error + + class TestMember(unittest.TestCase): def __init__(self, method_name='runTest'): @@ -125,49 +157,33 @@ class TestClient(unittest.TestCase): def set_up(self): socket.getaddrinfo = socket_getaddrinfo requests.get = requests_get - requests.put = requests_put - requests.delete = requests_delete dns.resolver.query = dns_query - self.client = Client({'discovery_srv': 'test'}) + with patch.object(etcd.Client, 'machines') as mock_machines: + mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001']) + self.client = Client({'discovery_srv': 'test'}) + self.client.http.request = http_request - def test__get(self): - self.assertRaises(EtcdError, self.client._get, '/bad_response') + def test_api_execute(self): + self.client._base_uri = 'http://localhost:4001' + self.client._machines_cache = ['http://localhost:2379'] + self.client.api_execute('/', 'GET') def test_get_srv_record(self): - self.assertEquals(Client.get_srv_record('blabla'), []) - self.assertEquals(Client.get_srv_record('exception'), []) + self.assertEquals(self.client.get_srv_record('blabla'), []) + self.assertEquals(self.client.get_srv_record('exception'), []) - def test_get_client_urls_from_dns(self): - self.assertEquals(Client.get_client_urls_from_dns('ok:2379'), ['http://127.0.0.1:2379/v2']) + def test__get_machines_cache_from_srv(self): + self.client.get_srv_record = lambda e: [('localhost', 2380)] + self.client._get_machines_cache_from_srv('blabla') - def test_load_members(self): - self.client._base_uri = self.client._base_uri.replace('localhost', 'error_code') - self.assertRaises(EtcdError, self.client.load_members) - self.client._base_uri = 'http://error_code:2380' - self.assertRaises(EtcdError, self.client.load_members) - self.client._base_uri = None + def test__get_machines_cache_from_dns(self): + self.client._get_machines_cache_from_dns('ok:2379') + + def test__load_machines_cache(self): self.client._config = {} - self.assertRaises(Exception, self.client.load_members) - - def test_get(self): - self.client._base_uri = None - self.assertRaises(EtcdConnectionFailed, self.client.get, '') - self.client._members_cache = ['http://error_code:4001/v2'] - self.client.get('') - - def test_put(self): - self.client._base_uri = None - self.assertRaises(EtcdConnectionFailed, self.client.put, '') - self.client._base_uri = 'http://localhost:4001/v2' - self.client._members_cache = ['http://error_code:4001/v2'] - self.client.put('') - - def test_delete(self): - self.client._base_uri = None - self.assertRaises(EtcdConnectionFailed, self.client.delete, '') - self.client._base_uri = 'http://localhost:4001/v2' - self.client._members_cache = ['http://error_code:4001/v2'] - self.client.delete('') + self.assertRaises(Exception, self.client._load_machines_cache) + self.client._config = {'discovery_srv': 'blabla'} + self.assertRaises(etcd.EtcdException, self.client._load_machines_cache) class TestEtcd(unittest.TestCase): @@ -177,37 +193,29 @@ class TestEtcd(unittest.TestCase): super(TestEtcd, self).__init__(method_name) def set_up(self): - socket.getaddrinfo = socket_getaddrinfo - requests.get = requests_get - requests.put = requests_put - requests.delete = requests_delete time.sleep = time_sleep - self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'}) + with patch.object(Client, 'machines') as mock_machines: + mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001']) + self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'}) + self.etcd.client.write = etcd_write + self.etcd.client.read = etcd_read def test_get_etcd_client(self): time.sleep = time_sleep_exception - self.assertRaises(Exception, self.etcd.get_etcd_client, {'host': 'error:2379'}) - - def test_get_client_path(self): - self.assertRaises(Exception, self.etcd.get_client_path, '', 2) - - def test_put_client_path(self): - self.assertRaises(EtcdError, self.etcd.put_client_path, '') - - def test_delete_client_path(self): - self.assertFalse(self.etcd.delete_client_path('')) + with patch.object(etcd.Client, 'machines') as mock_machines: + mock_machines.__get__ = Mock(side_effect=etcd.EtcdException) + self.assertRaises(Exception, self.etcd.get_etcd_client, {'discovery_srv': 'test'}) def test_get_cluster(self): - self.assertRaises(EtcdError, self.etcd.get_cluster) - self.etcd.client._base_uri = self.etcd.client._base_uri.replace('local', 'remote') + self.assertIsInstance(self.etcd.get_cluster(), Cluster) + self.etcd._base_path = '/service/nocluster' cluster = self.etcd.get_cluster() self.assertIsInstance(cluster, Cluster) - self.etcd.client._base_uri = self.etcd.client._base_uri.replace('remote', 'other') - self.etcd.get_cluster() - self.etcd.client._base_uri = self.etcd.client._base_uri.replace('other', 'noleader') - self.etcd.get_cluster() + self.assertIsNone(cluster.leader) def test_current_leader(self): + self.assertIsInstance(self.etcd.current_leader(), Member) + self.etcd._base_path = '/service/noleader' self.assertIsNone(self.etcd.current_leader()) def test_touch_member(self): @@ -216,14 +224,12 @@ class TestEtcd(unittest.TestCase): def test_take_leader(self): self.assertFalse(self.etcd.take_leader()) - def test_attempt_to_acquire_leader(self): - self.assertFalse(self.etcd.attempt_to_acquire_leader()) - def test_update_leader(self): - url = self.etcd.client._base_uri = self.etcd.client._base_uri.replace('local', 'remote') self.assertTrue(self.etcd.update_leader(MockPostgresql())) - self.etcd.client._base_uri = url.replace('remote', 'other') - self.assertFalse(self.etcd.update_leader(MockPostgresql())) def test_race(self): self.assertFalse(self.etcd.race('')) + + def test_delete_leader(self): + self.etcd.client.delete = etcd_delete + self.assertFalse(self.etcd.delete_leader()) diff --git a/tests/test_ha.py b/tests/test_ha.py index be7ffcda..abfc4ac8 100644 --- a/tests/test_ha.py +++ b/tests/test_ha.py @@ -1,10 +1,10 @@ import unittest -import requests -from helpers.dcs import DCSError -from helpers.etcd import Cluster, Etcd +from helpers.dcs import Cluster, DCSError +from helpers.etcd import Client, Etcd from helpers.ha import Ha -from test_etcd import requests_get, requests_put, requests_delete +from mock import Mock, patch +from test_etcd import etcd_read, etcd_write def true(*args, **kwargs): @@ -71,15 +71,16 @@ class TestHa(unittest.TestCase): super(TestHa, self).__init__(method_name) def set_up(self): - requests.get = requests_get - requests.put = requests_put - requests.delete = requests_delete self.p = MockPostgresql() - self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'}) - self.ha = Ha(self.p, self.e) - self.ha.load_cluster_from_dcs() - self.ha.cluster = get_unlocked_cluster() - self.ha.load_cluster_from_dcs = nop + with patch.object(Client, 'machines') as mock_machines: + mock_machines.__get__ = Mock(return_value=['http://remotehost:2379']) + self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'}) + self.e.client.read = etcd_read + self.e.client.write = etcd_write + self.ha = Ha(self.p, self.e) + self.ha.load_cluster_from_dcs() + self.ha.cluster = get_unlocked_cluster() + self.ha.load_cluster_from_dcs = nop def test_load_cluster_from_dcs(self): ha = Ha(self.p, self.e) diff --git a/tests/test_patroni.py b/tests/test_patroni.py index a2d29e9c..7d772f68 100644 --- a/tests/test_patroni.py +++ b/tests/test_patroni.py @@ -1,18 +1,20 @@ import datetime import helpers.zookeeper import psycopg2 -import requests import subprocess import sys import time import unittest import yaml -from patroni import Patroni, main +from helpers.api import RestApiServer from helpers.dcs import Cluster, Member +from helpers.etcd import Etcd from helpers.zookeeper import ZooKeeper +from mock import Mock, patch +from patroni import Patroni, main from six.moves import BaseHTTPServer -from test_etcd import requests_get, requests_put, requests_delete +from test_etcd import Client, etcd_read, etcd_write from test_ha import true, false from test_postgresql import Postgresql, subprocess_call, psycopg2_connect from test_zookeeper import MockKazooClient @@ -26,6 +28,15 @@ def time_sleep(*args): raise Exception() +class Mock_BaseServer__is_shut_down: + + def set(self): + pass + + def clear(self): + pass + + class TestPatroni(unittest.TestCase): def __init__(self, method_name='runTest'): @@ -37,9 +48,6 @@ class TestPatroni(unittest.TestCase): self.touched = False subprocess.call = subprocess_call psycopg2.connect = psycopg2_connect - requests.get = requests_get - requests.put = requests_put - requests.delete = requests_delete self.time_sleep = time.sleep time.sleep = nop self.write_pg_hba = Postgresql.write_pg_hba @@ -47,9 +55,13 @@ class TestPatroni(unittest.TestCase): Postgresql.write_pg_hba = nop Postgresql.write_recovery_conf = nop BaseHTTPServer.HTTPServer.__init__ = nop + RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down() + RestApiServer._BaseServer__shutdown_request = True with open('postgres0.yml', 'r') as f: config = yaml.load(f) - self.g = Patroni(config) + with patch.object(Client, 'machines') as mock_machines: + mock_machines.__get__ = Mock(return_value=['http://remotehost:2379']) + self.p = Patroni(config) def tear_down(self): time.sleep = self.time_sleep @@ -58,50 +70,74 @@ class TestPatroni(unittest.TestCase): def test_get_dcs(self): helpers.zookeeper.KazooClient = MockKazooClient - self.assertIsInstance(self.g.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper) - self.assertRaises(Exception, self.g.get_dcs, '', {}) + self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper) + self.assertRaises(Exception, self.p.get_dcs, '', {}) def test_patroni_main(self): main() sys.argv = ['patroni.py', 'postgres0.yml'] time.sleep = time_sleep - self.assertRaises(Exception, main) + + with patch.object(Client, 'machines') as mock_machines: + mock_machines.__get__ = Mock(return_value=['http://remotehost:2379']) + Patroni.initialize = nop + touch_member = Patroni.touch_member + run = Patroni.run + + Patroni.touch_member = self.touch_member + Patroni.run = time_sleep + + Etcd.delete_leader = nop + + self.assertRaises(Exception, main) + + Patroni.run = run + Patroni.touch_member = touch_member def test_patroni_run(self): time.sleep = time_sleep - self.g.postgresql.is_leader = lambda: False - self.g.ha.state_handler.sync_replication_slots = time_sleep - self.assertRaises(Exception, self.g.run) + self.p.touch_member = self.touch_member + self.p.ha.state_handler.sync_replication_slots = time_sleep + self.p.ha.dcs.client.read = etcd_read + self.assertRaises(Exception, self.p.run) + self.p.ha.state_handler.is_leader = lambda: False + self.p.api.start = nop + self.assertRaises(Exception, self.p.run) - def touch_member(self): + def touch_member(self, ttl=None): if not self.touched: self.touched = True return False return True def test_touch_member(self): + self.p.ha.dcs.client.write = etcd_write + self.p.touch_member() now = datetime.datetime.utcnow() - member = Member(0, self.g.postgresql.name, 'b', 'c', (now + datetime.timedelta( - seconds=self.g.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None) - self.g.ha.cluster = Cluster(True, member, 0, [member]) - self.g.touch_member() + member = Member(0, self.p.postgresql.name, 'b', 'c', (now + datetime.timedelta( + seconds=self.p.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None) + self.p.ha.cluster = Cluster(True, member, 0, [member]) + self.p.touch_member() def test_patroni_initialize(self): - self.g.postgresql.should_use_s3_to_create_replica = false - self.g.ha.dcs.client._base_uri = 'http://remote' - self.g.postgresql.data_directory_empty = true - self.g.ha.dcs.race = true - self.g.initialize() - self.g.ha.dcs.race = false - self.g.initialize() - self.g.postgresql.data_directory_empty = false - self.g.touch_member = self.touch_member - self.g.initialize() - self.g.postgresql.data_directory_empty = true + self.p.postgresql.should_use_s3_to_create_replica = false + self.p.ha.dcs.client.write = etcd_write + self.p.touch_member = self.touch_member + self.p.postgresql.data_directory_empty = true + self.p.ha.dcs.race = true + self.p.initialize() + + self.p.ha.dcs.race = false time.sleep = time_sleep - self.g.postgresql.sync_from_leader = false - self.assertRaises(Exception, self.g.initialize) + self.p.ha.dcs.client.read = etcd_read + self.p.initialize() + + self.p.ha.dcs.current_leader = nop + self.assertRaises(Exception, self.p.initialize) + + self.p.postgresql.data_directory_empty = false + self.p.initialize() def test_schedule_next_run(self): - self.g.next_run = time.time() - self.g.nap_time - 1 - self.g.schedule_next_run() + self.p.next_run = time.time() - self.p.nap_time - 1 + self.p.schedule_next_run() diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 6cbb6cb7..68ba47bc 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -1,6 +1,3 @@ -#!/usr/bin/env python -# -*- coding: utf-8 -*- - import os import psycopg2 import shutil diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py index 172a8c5f..53f5ab82 100644 --- a/tests/test_zookeeper.py +++ b/tests/test_zookeeper.py @@ -1,6 +1,6 @@ import helpers.zookeeper -import unittest import requests +import unittest from helpers.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError from kazoo.client import KazooState