patroni/tests/test_zookeeper.py

import helpers.zookeeper
import requests
import unittest

from helpers.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


class MockEvent:

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

    def __init__(self, **kwargs):
        self.handler = MockEventHandler()
        self.leader = False
        self.exists = True

    def start(self, timeout):
        pass

    @property
    def client_id(self):
        return (-1, '')

    def add_listener(self, cb):
        pass

    def retry(self, func, *args, **kwargs):
        func(*args, **kwargs)

    def get(self, path, watch=None):
        if path == '/service/test/no_node':
            raise NoNodeError
        elif path == '/service/test/other_exception':
            raise Exception()
        elif '/members/' in path:
            return (
                'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni',
                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
            )
        elif path.endswith('/optime/leader'):
            return '1'
        elif path.endswith('/leader'):
            if self.leader:
                return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
            return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))

    def get_children(self, path, watch=None, include_data=False):
        return ['foo', 'bar', 'buzz']

    def create(self, path, value="", acl=None, ephemeral=False, sequence=False, makepath=False):
        if path.endswith('/initialize') or path == '/service/test/optime/leader':
            raise Exception
        elif value == 'retry' or (value == 'exists' and self.exists):
            raise NodeExistsError

    def set(self, path, value, version=-1):
        if path == '/service/bla/optime/leader':
            raise Exception
        raise NoNodeError

    def delete(self, path, version=-1, recursive=False):
        self.exists = False
        if path == '/service/test/leader':
            if self.leader:
                return
            self.leader = True
            raise Exception

    def set_hosts(self, hosts, randomize_hosts=None):
        pass


def exhibitor_sleep(_):
    raise Exception


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
        helpers.zookeeper.sleep = exhibitor_sleep

    def test_init(self):
        self.assertRaises(Exception, 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
        helpers.zookeeper.KazooClient = MockKazooClient
        self.zk = ZooKeeper('foo', {'exhibitor': {'hosts': ['localhost', 'exhibitor'], 'port': 8181}, 'scope': 'test'})

    def test_session_listener(self):
        self.zk.session_listener(KazooState.SUSPENDED)

    def test_get_node(self):
        self.assertIsNone(self.zk.get_node('/no_node'))
        self.assertIsNone(self.zk.get_node('/other_exception'))

    def test__inner_load_cluster(self):
        self.zk._base_path = self.zk._base_path.replace('test', 'bla')
        self.zk._inner_load_cluster()

    def test_get_cluster(self):
        self.assertRaises(ZooKeeperError, self.zk.get_cluster)
        self.zk.exhibitor.poll = lambda: True
        self.zk.get_cluster()
        self.zk.touch_member('foo')
        self.zk.delete_leader()

    def test_race(self):
        self.assertFalse(self.zk.initialize())

    def test_touch_member(self):
        self.zk.touch_member('new')
        self.zk.touch_member('exists')
        self.zk.touch_member('retry')

    def test_take_leader(self):
        self.zk.take_leader()

    def test_update_leader(self):
        self.zk.last_leader_operation = -1
        self.assertTrue(self.zk.update_leader(MockPostgresql()))
        self.zk._base_path = self.zk._base_path.replace('test', 'bla')
        self.zk.last_leader_operation = -1
        self.assertTrue(self.zk.update_leader(MockPostgresql()))

    def test_sleep(self):
        self.zk.sleep(0)