Repository: zalando/patroni
Pull request #2909 (2023-10-12 11:56:44)

Don't cache anything in Zookeeper implementation
Cache creates a lot of problems and prevents implementing a feature of automatic retention of physical replication slots for members with configurable retention policy.

Just read the entire cluster from Zookeeper instead and use watchers only for the `/leader` and `/config` keys.

coveralls: ## Pull Request Test Coverage Report for [Build 6495276291](https://coveralls.io/builds/63249638)

* **34** of **34**   **(100.0%)**  changed or added relevant lines in **3** files are covered.
* No unchanged relevant lines lost coverage.
* Overall coverage remained the same at **99.863%**

---



|  Totals | [![Coverage Status](https://coveralls.io/builds/63249638/badge)](https://coveralls.io/builds/63249638) |
| :-- | --: |
| Change from base [Build 6481756789](https://coveralls.io/builds/63220042): |  0.0% |
| Covered Lines: | 13100 |
| Relevant Lines: | 13118 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/.github/workflows/install_deps.py b/.github/workflows/install_deps.py
index ea0c9d1b..6480f66a 100644
--- a/.github/workflows/install_deps.py
+++ b/.github/workflows/install_deps.py
@@ -45,8 +45,8 @@ def install_packages(what):
     packages['exhibitor'] = packages['zookeeper']
     packages = packages.get(what, [])
     ver = versions.get(what)
-    if float(ver) == 15:
-        packages += ['postgresql-{0}-citus-12.0'.format(ver)]
+    if float(ver) >= 15:
+        packages += ['postgresql-{0}-citus-12.1'.format(ver)]
     subprocess.call(['sudo', 'apt-get', 'update', '-y'])
     return subprocess.call(['sudo', 'apt-get', 'install', '-y', 'postgresql-' + ver, 'expect-dev'] + packages)
 
diff --git a/docker/README.md b/docker/README.md
index f2b30ab6..da87842f 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -78,7 +78,7 @@ Example session:
     |   demo  | patroni3 | 172.22.0.4 |        | running |  1 |         0 |
     +---------+----------+------------+--------+---------+----+-----------+
 
-    postgres@patroni1:~$ etcdctl ls --recursive --sort -p /service/demo
+    postgres@patroni1:~$ etcdctl get --keys-only --prefix /service/demo
     /service/demo/config
     /service/demo/initialize
     /service/demo/leader
diff --git a/docs/citus.rst b/docs/citus.rst
index 084931df..62cbdd1b 100644
--- a/docs/citus.rst
+++ b/docs/citus.rst
@@ -38,14 +38,18 @@ After that you just need to start Patroni and it will handle the rest:
 2. If ``max_prepared_transactions`` isn't explicitly set in the global
    :ref:`dynamic configuration <dynamic_configuration>` Patroni will
    automatically set it to ``2*max_connections``.
-3. The ``citus.database`` will be automatically created followed by ``CREATE EXTENSION citus``.
-4. Current superuser :ref:`credentials <postgresql_settings>` will be added to the ``pg_dist_authinfo``
+3. The ``citus.local_hostname`` GUC value will be adjusted from ``localhost`` to the
+   value that Patroni is using in order to connect to the local PostgreSQL
+   instance. The value sometimes should be different from the ``localhost``
+   because PostgreSQL might be not listening on it.
+4. The ``citus.database`` will be automatically created followed by ``CREATE EXTENSION citus``.
+5. Current superuser :ref:`credentials <postgresql_settings>` will be added to the ``pg_dist_authinfo``
    table to allow cross-node communication. Don't forget to update them if
    later you decide to change superuser username/password/sslcert/sslkey!
-5. The coordinator primary node will automatically discover worker primary
+6. The coordinator primary node will automatically discover worker primary
    nodes and add them to the ``pg_dist_node`` table using the
    ``citus_add_node()`` function.
-6. Patroni will also maintain ``pg_dist_node`` in case failover/switchover
+7. Patroni will also maintain ``pg_dist_node`` in case failover/switchover
    on the coordinator or worker clusters occurs.
 
 patronictl
diff --git a/features/citus.feature b/features/citus.feature
index 35ccebbe..b23eb2b2 100644
--- a/features/citus.feature
+++ b/features/citus.feature
@@ -68,6 +68,6 @@ Feature: citus
     And I receive a response output "+ttl: 20"
     Then postgres4 is registered in the postgres2 as the primary in group 2 after 5 seconds
     When I shut down postgres4
-    Then There is a transaction in progress on postgres0 changing pg_dist_node
+    Then there is a transaction in progress on postgres0 changing pg_dist_node after 5 seconds
     When I run patronictl.py restart batman postgres2 --group 1 --force
     Then a transaction finishes in 20 seconds
diff --git a/features/environment.py b/features/environment.py
index 79db5d16..3e0ad3b1 100644
--- a/features/environment.py
+++ b/features/environment.py
@@ -245,6 +245,10 @@ class PatroniController(AbstractController):
             self.recursive_update(config, custom_config)
 
         self.recursive_update(config, {
+            'log': {
+                'format': '%(asctime)s %(levelname)s [%(pathname)s:%(lineno)d - %(funcName)s]: %(message)s',
+                'loggers': {'patroni.postgresql.callback_executor': 'DEBUG'}
+            },
             'bootstrap': {
                 'dcs': {
                     'loop_wait': 2,
@@ -688,7 +692,7 @@ class ZooKeeperController(AbstractExternalDcsController):
         self._client = kazoo.client.KazooClient()
 
     def process_name(self):
-        return "zookeeper"
+        return "java .*zookeeper"
 
     def query(self, key, scope='batman', group=None):
         import kazoo.exceptions
diff --git a/features/steps/citus.py b/features/steps/citus.py
index 644219c7..4dd2ffa6 100644
--- a/features/steps/citus.py
+++ b/features/steps/citus.py
@@ -115,12 +115,18 @@ def count_rows(context, name):
     assert rows == context.insert_counter, "Distributed table doesn't have expected amount of rows"
 
 
-@step("There is a transaction in progress on {name:w} changing pg_dist_node")
-def check_transaction(context, name):
-    cur = context.pctl.query(name, "SELECT xact_start FROM pg_stat_activity WHERE pid <> pg_backend_pid()"
-                                   " AND state = 'idle in transaction' AND query ~ 'citus_update_node'")
-    assert cur.rowcount == 1, "There is no idle in transaction updating pg_dist_node"
-    context.xact_start = cur.fetchone()[0]
+@step("there is a transaction in progress on {name:w} changing pg_dist_node after {time_limit:d} seconds")
+def check_transaction(context, name, time_limit):
+    time_limit *= context.timeout_multiplier
+    max_time = time.time() + int(time_limit)
+    while time.time() < max_time:
+        cur = context.pctl.query(name, "SELECT xact_start FROM pg_stat_activity WHERE pid <> pg_backend_pid()"
+                                       " AND state = 'idle in transaction' AND query ~ 'citus_update_node'")
+        if cur.rowcount == 1:
+            context.xact_start = cur.fetchone()[0]
+            return
+        time.sleep(1)
+    assert False, f"There is no idle in transaction on {name} updating pg_dist_node after {time_limit} seconds"
 
 
 @step("a transaction finishes in {timeout:d} seconds")
diff --git a/patroni/api.py b/patroni/api.py
index 911ea574..e1b71337 100644
--- a/patroni/api.py
+++ b/patroni/api.py
@@ -451,7 +451,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
         Write an HTTP response with JSON content based on the output of :func:`~patroni.utils.cluster_as_json`, with
         HTTP status ``200`` and the JSON representation of the cluster topology.
         """
-        cluster = self.server.patroni.dcs.get_cluster(True)
+        cluster = self.server.patroni.dcs.get_cluster()
         global_config = self.server.patroni.config.get_global_config(cluster)
 
         response = cluster_as_json(cluster, global_config)
@@ -690,7 +690,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
         """
         request = self._read_json_content()
         if request:
-            cluster = self.server.patroni.dcs.get_cluster(True)
+            cluster = self.server.patroni.dcs.get_cluster()
             if not (cluster.config and cluster.config.modify_version):
                 return self.send_error(503)
             data = cluster.config.data.copy()
@@ -1166,7 +1166,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
 
         patroni = self.server.patroni
         if patroni.postgresql.citus_handler.is_coordinator() and patroni.ha.is_leader():
-            cluster = patroni.dcs.get_cluster(True)
+            cluster = patroni.dcs.get_cluster()
             patroni.postgresql.citus_handler.handle_event(cluster, request)
         self.write_response(200, 'OK')
 
diff --git a/patroni/dcs/__init__.py b/patroni/dcs/__init__.py
index d71d41ad..1a516cca 100644
--- a/patroni/dcs/__init__.py
+++ b/patroni/dcs/__init__.py
@@ -1548,9 +1548,6 @@ class AbstractDCS(abc.ABC):
                 primary and exception raised, instance would be demoted.
         """
 
-    def _bypass_caches(self) -> None:
-        """Used only in Zookeeper."""
-
     def __get_patroni_cluster(self, path: Optional[str] = None) -> Cluster:
         """Low level method to load a :class:`Cluster` object from DCS.
 
@@ -1593,14 +1590,13 @@ class AbstractDCS(abc.ABC):
                   dict.
         """
         groups = self._load_cluster(self._base_path + '/', self._citus_cluster_loader)
-        if isinstance(groups, Cluster):  # Zookeeper could return a cached version
-            cluster = groups
-        else:
-            cluster = groups.pop(CITUS_COORDINATOR_GROUP_ID, Cluster.empty())
-            cluster.workers.update(groups)
+        if TYPE_CHECKING:  # pragma: no cover
+            assert isinstance(groups, dict)
+        cluster = groups.pop(CITUS_COORDINATOR_GROUP_ID, Cluster.empty())
+        cluster.workers.update(groups)
         return cluster
 
-    def get_cluster(self, force: bool = False) -> Cluster:
+    def get_cluster(self) -> Cluster:
         """Retrieve an appropriate cached or fresh view of DCS.
 
         .. note::
@@ -1609,12 +1605,8 @@ class AbstractDCS(abc.ABC):
 
             Returns either a Citus or Patroni implementation of :class:`Cluster` depending on availability.
 
-        :param force: a value of ``True`` will override Zookeeper caching features.
-
         :returns:
         """
-        if force:
-            self._bypass_caches()
         try:
             cluster = self._get_citus_cluster() if self.is_citus_coordinator() else self.__get_patroni_cluster()
         except Exception:
diff --git a/patroni/dcs/zookeeper.py b/patroni/dcs/zookeeper.py
index 649c1ac5..3093ba06 100644
--- a/patroni/dcs/zookeeper.py
+++ b/patroni/dcs/zookeeper.py
@@ -116,10 +116,7 @@ class ZooKeeper(AbstractDCS):
                                           timeout=config['ttl'], connection_retry=KazooRetry(max_delay=1, max_tries=-1,
                                           sleep_func=time.sleep), command_retry=KazooRetry(max_delay=1, max_tries=-1,
                                           deadline=config['retry_timeout'], sleep_func=time.sleep), **kwargs)
-        self._client.add_listener(self.session_listener)
 
-        self._fetch_cluster: bool = True
-        self._fetch_status: bool = True
         self.__last_member_data: Optional[Dict[str, Any]] = None
 
         self._orig_kazoo_connect = self._client._connection._connect
@@ -142,18 +139,9 @@ class ZooKeeper(AbstractDCS):
         ret = self._orig_kazoo_connect(*args)
         return max(self.loop_wait - 2, 2) * 1000, ret[1]
 
-    def session_listener(self, state: str) -> None:
-        if state in [KazooState.SUSPENDED, KazooState.LOST]:
-            self.cluster_watcher(None)
-
-    def status_watcher(self, event: Optional[WatchedEvent]) -> None:
-        self._fetch_status = True
-        self.event.set()
-
-    def cluster_watcher(self, event: Optional[WatchedEvent]) -> None:
-        self._fetch_cluster = True
-        if not event or event.state != KazooState.CONNECTED or event.path.startswith(self.client_path('')):
-            self.status_watcher(event)
+    def _watcher(self, event: WatchedEvent) -> None:
+        if event.state != KazooState.CONNECTED or event.path.startswith(self.client_path('')):
+            self.event.set()
 
     def reload_config(self, config: Union['Config', Dict[str, Any]]) -> None:
         self.set_retry_timeout(config['retry_timeout'])
@@ -202,75 +190,66 @@ class ZooKeeper(AbstractDCS):
             return None
 
     def get_status(self, path: str, leader: Optional[Leader]) -> Status:
-        watch = self.status_watcher if not leader or leader.name != self._name else None
-
-        status = self.get_node(path + self._STATUS, watch)
+        status = self.get_node(path + self._STATUS)
         if not status:
-            status = self.get_node(path + self._LEADER_OPTIME, watch)
-        if status:
-            self._fetch_status = False
+            status = self.get_node(path + self._LEADER_OPTIME)
         return Status.from_node(status and status[0])
 
     @staticmethod
     def member(name: str, value: str, znode: ZnodeStat) -> Member:
         return Member.from_node(znode.version, name, znode.ephemeralOwner, value)
 
-    def get_children(self, key: str, watch: Optional[Callable[[WatchedEvent], None]] = None) -> List[str]:
+    def get_children(self, key: str) -> List[str]:
         try:
-            return self._client.get_children(key, watch)
+            return self._client.get_children(key)
         except NoNodeError:
             return []
 
     def load_members(self, path: str) -> List[Member]:
         members: List[Member] = []
-        for member in self.get_children(path + self._MEMBERS, self.cluster_watcher):
+        for member in self.get_children(path + self._MEMBERS):
             data = self.get_node(path + self._MEMBERS + member)
             if data is not None:
                 members.append(self.member(member, *data))
         return members
 
     def _cluster_loader(self, path: str) -> Cluster:
-        self._fetch_cluster = False
-        self.event.clear()
-        nodes = set(self.get_children(path, self.cluster_watcher))
-        if not nodes:
-            self._fetch_cluster = True
+        nodes = set(self.get_children(path))
 
         # get initialize flag
         initialize = (self.get_node(path + self._INITIALIZE) or [None])[0] if self._INITIALIZE in nodes else None
 
         # get global dynamic configuration
-        config = self.get_node(path + self._CONFIG, watch=self.cluster_watcher) if self._CONFIG in nodes else None
+        config = self.get_node(path + self._CONFIG, watch=self._watcher) if self._CONFIG in nodes else None
         config = config and ClusterConfig.from_node(config[1].version, config[0], config[1].mzxid)
 
         # get timeline history
-        history = self.get_node(path + self._HISTORY, watch=self.cluster_watcher) if self._HISTORY in nodes else None
+        history = self.get_node(path + self._HISTORY) if self._HISTORY in nodes else None
         history = history and TimelineHistory.from_node(history[1].mzxid, history[0])
 
         # get synchronization state
-        sync = self.get_node(path + self._SYNC, watch=self.cluster_watcher) if self._SYNC in nodes else None
+        sync = self.get_node(path + self._SYNC) if self._SYNC in nodes else None
         sync = SyncState.from_node(sync and sync[1].version, sync and sync[0])
 
         # get list of members
         members = self.load_members(path) if self._MEMBERS[:-1] in nodes else []
 
         # get leader
-        leader = self.get_node(path + self._LEADER) if self._LEADER in nodes else None
+        leader = self.get_node(path + self._LEADER, watch=self._watcher) if self._LEADER in nodes else None
         if leader:
             member = Member(-1, leader[0], None, {})
             member = ([m for m in members if m.name == leader[0]] or [member])[0]
             leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
-            self._fetch_cluster = member.version == -1
 
         # get last known leader lsn and slots
         status = self.get_status(path, leader)
 
         # failover key
-        failover = self.get_node(path + self._FAILOVER, watch=self.cluster_watcher) if self._FAILOVER in nodes else None
+        failover = self.get_node(path + self._FAILOVER) if self._FAILOVER in nodes else None
         failover = failover and Failover.from_node(failover[1].version, failover[0])
 
         # get failsafe topology
-        failsafe = self.get_node(path + self._FAILSAFE, watch=self.cluster_watcher) if self._FAILSAFE in nodes else None
+        failsafe = self.get_node(path + self._FAILSAFE) if self._FAILSAFE in nodes else None
         try:
             failsafe = json.loads(failsafe[0]) if failsafe else None
         except Exception:
@@ -279,45 +258,20 @@ class ZooKeeper(AbstractDCS):
         return Cluster(initialize, config, leader, status, members, failover, sync, history, failsafe)
 
     def _citus_cluster_loader(self, path: str) -> Dict[int, Cluster]:
-        fetch_cluster = False
         ret: Dict[int, Cluster] = {}
-        for node in self.get_children(path, self.cluster_watcher):
+        for node in self.get_children(path):
             if citus_group_re.match(node):
                 ret[int(node)] = self._cluster_loader(path + node + '/')
-                fetch_cluster = fetch_cluster or self._fetch_cluster
-        self._fetch_cluster = fetch_cluster
         return ret
 
     def _load_cluster(
             self, path: str, loader: Callable[[str], Union[Cluster, Dict[int, Cluster]]]
     ) -> Union[Cluster, Dict[int, Cluster]]:
-        cluster = self.cluster if path == self._base_path + '/' else None
-        if self._fetch_cluster or cluster is None:
-            try:
-                cluster = self._client.retry(loader, path)
-            except Exception:
-                logger.exception('get_cluster')
-                self.cluster_watcher(None)
-                raise ZooKeeperError('ZooKeeper in not responding properly')
-        # The /status ZNode was updated or doesn't exist
-        elif self._fetch_status and not self._fetch_cluster or not cluster.last_lsn \
-                or cluster.has_permanent_slots(self._name) and not cluster.slots:
-            # If current node is the leader just clear the event without fetching anything (we are updating the /status)
-            if cluster.leader and cluster.leader.name == self._name:
-                self.event.clear()
-            else:
-                try:
-                    status = self.get_status(self.client_path(''), cluster.leader)
-                    self.event.clear()
-                    new_cluster: List[Any] = list(cluster)
-                    new_cluster[3] = status
-                    cluster = Cluster(*new_cluster)
-                except Exception:
-                    pass
-        return cluster
-
-    def _bypass_caches(self) -> None:
-        self._fetch_cluster = True
+        try:
+            return self._client.retry(loader, path)
+        except Exception:
+            logger.exception('get_cluster')
+            raise ZooKeeperError('ZooKeeper in not responding properly')
 
     def _create(self, path: str, value: bytes, retry: bool = False, ephemeral: bool = False) -> bool:
         try:
@@ -334,7 +288,6 @@ class ZooKeeper(AbstractDCS):
         try:
             self._client.retry(self._client.create, self.leader_path, self._name.encode('utf-8'),
                                makepath=True, ephemeral=True)
-            self.cluster_watcher(None)  # the next _load_cluster() call must read from ZooKeeper.
             return True
         except (ConnectionClosedError, RetryFailedError) as e:
             raise ZooKeeperError(e)
@@ -381,20 +334,9 @@ class ZooKeeper(AbstractDCS):
         member = cluster and cluster.get_member(self._name, fallback_to_leader=False)
         member_data = self.__last_member_data or member and member.data
         if member and member_data:
-            is_leader = data.get('role') in ('master', 'primary', 'standby_leader')
-            checkpoint_after_promote_changed = member_data.get('checkpoint_after_promote') \
-                != data.get('checkpoint_after_promote')
-            state_running_changed = member_data.get('state') != data.get('state') \
-                and 'running' in (member_data.get('state'), data.get('state'))
-            tags_changed = not deep_compare(member_data.get('tags', {}), data.get('tags', {}))
-
-            # We want delete the member ZNode if:
-            # - our session doesn't match with session id on our member key; or
-            # - we want to notify leader if some important fields in the member key changed; or
-            # - if we are the leader and want to notify replicas about checkpoint_after_promote;
-            if self._client.client_id is not None and member.session != self._client.client_id[0] \
-                    or is_leader and checkpoint_after_promote_changed \
-                    or not is_leader and (state_running_changed or tags_changed):
+            # We want delete the member ZNode if our session doesn't match with session id on our member key
+            if self._client.client_id is not None and member.session != self._client.client_id[0]:
+                logger.warning('Recreating the member ZNode due to ownership mismatch')
                 try:
                     self._client.delete_async(self.member_path).get(timeout=1)
                 except NoNodeError:
@@ -493,7 +435,10 @@ class ZooKeeper(AbstractDCS):
         return self.set_sync_state_value("{}", version) is not False
 
     def watch(self, leader_version: Optional[int], timeout: float) -> bool:
-        ret = super(ZooKeeper, self).watch(leader_version, timeout + 0.5)
-        if ret and not self._fetch_status:
-            self._fetch_cluster = True
-        return ret or self._fetch_cluster
+        if leader_version:
+            timeout += 0.5
+
+        try:
+            return super(ZooKeeper, self).watch(leader_version, timeout)
+        finally:
+            self.event.clear()
diff --git a/patroni/postgresql/callback_executor.py b/patroni/postgresql/callback_executor.py
index fa645b86..06b9f353 100644
--- a/patroni/postgresql/callback_executor.py
+++ b/patroni/postgresql/callback_executor.py
@@ -1,8 +1,9 @@
 import logging
+import sys
 
 from enum import Enum
 from threading import Condition, Thread
-from typing import List
+from typing import Any, Dict, List
 
 from .cancellable import CancellableExecutor, CancellableSubprocess
 
@@ -53,6 +54,8 @@ class CallbackExecutor(CancellableExecutor, Thread):
         If it couldn't be killed we wait until it finishes.
 
         :param cmd: command to be executed"""
+        kwargs: Dict[str, Any] = {'stacklevel': 3} if sys.version_info >= (3, 8) else {}
+        logger.debug('CallbackExecutor.call(%s)', cmd, **kwargs)
 
         if cmd[-3] == CallbackAction.ON_RELOAD:
             return self._on_reload_executor.call_nowait(cmd)
diff --git a/patroni/postgresql/citus.py b/patroni/postgresql/citus.py
index 8ca2790e..26923f37 100644
--- a/patroni/postgresql/citus.py
+++ b/patroni/postgresql/citus.py
@@ -403,6 +403,9 @@ class CitusHandler(Thread):
         # Resharding in Citus implemented using logical replication
         parameters['wal_level'] = 'logical'
 
+        # Sometimes Citus needs to connect to the local postgres. We will do it the same way as Patroni does.
+        parameters['citus.local_hostname'] = self._postgresql.connection_pool.conn_kwargs.get('host', 'localhost')
+
     def ignore_replication_slot(self, slot: Dict[str, str]) -> bool:
         if isinstance(self._config, dict) and self._postgresql.is_primary() and\
                 slot['type'] == 'logical' and slot['database'] == self._config['database']:


Test patch:
diff --git a/tests/__init__.py b/tests/__init__.py
index 32172926..bd70ba3d 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -241,7 +241,7 @@ class PostgresInit(unittest.TestCase):
                                                 'replication': {'username': '', 'password': 'rep-pass'},
                                                 'rewind': {'username': 'rewind', 'password': 'test'}},
                              'remove_data_directory_on_rewind_failure': True,
-                             'use_pg_rewind': True, 'pg_ctl_timeout': 'bla',
+                             'use_pg_rewind': True, 'pg_ctl_timeout': 'bla', 'use_unix_socket': True,
                              'parameters': self._PARAMETERS,
                              'recovery_conf': {'foo': 'bar'},
                              'pg_hba': ['host all all 0.0.0.0/0 md5'],
diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py
index 9ac6aa68..4c2d1c98 100644
--- a/tests/test_bootstrap.py
+++ b/tests/test_bootstrap.py
@@ -256,7 +256,7 @@ class TestBootstrap(BaseTestPostgresql):
         mock_cancellable_subprocess_call.assert_called()
         args, kwargs = mock_cancellable_subprocess_call.call_args
         self.assertTrue('PGPASSFILE' in kwargs['env'])
-        self.assertEqual(args[0], ['/bin/false', 'dbname=postgres host=127.0.0.2 port=5432'])
+        self.assertEqual(args[0], ['/bin/false', 'dbname=postgres host=/tmp port=5432'])
 
         mock_cancellable_subprocess_call.reset_mock()
         self.p.connection_pool._conn_kwargs.pop('host')
diff --git a/tests/test_citus.py b/tests/test_citus.py
index 9849a069..7279893e 100644
--- a/tests/test_citus.py
+++ b/tests/test_citus.py
@@ -13,7 +13,6 @@ class TestCitus(BaseTestPostgresql):
     def setUp(self):
         super(TestCitus, self).setUp()
         self.c = self.p.citus_handler
-        self.p.connection_pool.conn_kwargs = {'host': 'localhost', 'dbname': 'postgres'}
         self.cluster = get_cluster_initialized_with_leader()
         self.cluster.workers[1] = self.cluster
 
@@ -139,6 +138,7 @@ class TestCitus(BaseTestPostgresql):
         self.assertEqual(parameters['max_prepared_transactions'], 202)
         self.assertEqual(parameters['shared_preload_libraries'], 'citus,foo,bar')
         self.assertEqual(parameters['wal_level'], 'logical')
+        self.assertEqual(parameters['citus.local_hostname'], '/tmp')
 
     def test_bootstrap(self):
         self.c._config = None
diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py
index 184fda05..3ce3ea75 100644
--- a/tests/test_zookeeper.py
+++ b/tests/test_zookeeper.py
@@ -1,13 +1,13 @@
 import select
 import unittest
 
-from kazoo.client import KazooClient, KazooState
+from kazoo.client import KazooClient
 from kazoo.exceptions import NoNodeError, NodeExistsError
 from kazoo.handlers.threading import SequentialThreadingHandler
-from kazoo.protocol.states import KeeperState, ZnodeStat
+from kazoo.protocol.states import KeeperState, WatchedEvent, ZnodeStat
 from kazoo.retry import RetryFailedError
 from mock import Mock, PropertyMock, patch
-from patroni.dcs.zookeeper import Cluster, Leader, PatroniKazooClient, \
+from patroni.dcs.zookeeper import Cluster, PatroniKazooClient, \
     PatroniSequentialThreadingHandler, ZooKeeper, ZooKeeperError
 
 
@@ -152,9 +152,6 @@ class TestZooKeeper(unittest.TestCase):
                              'name': 'foo', 'ttl': 30, 'retry_timeout': 10, 'loop_wait': 10,
                              'set_acls': {'CN=principal2': ['ALL']}})
 
-    def test_session_listener(self):
-        self.zk.session_listener(KazooState.SUSPENDED)
-
     def test_reload_config(self):
         self.zk.reload_config({'ttl': 20, 'retry_timeout': 10, 'loop_wait': 10})
         self.zk.reload_config({'ttl': 20, 'retry_timeout': 10, 'loop_wait': 5})
@@ -176,15 +173,6 @@ class TestZooKeeper(unittest.TestCase):
         self.zk._cluster_loader(self.zk.client_path(''))
 
     def test_get_cluster(self):
-        cluster = self.zk.get_cluster(True)
-        self.assertIsInstance(cluster.leader, Leader)
-        self.zk.status_watcher(None)
-        self.zk.get_cluster()
-        self.zk.touch_member({'foo': 'foo'})
-        self.zk._name = 'bar'
-        self.zk.status_watcher(None)
-        with patch.object(ZooKeeper, 'get_node', Mock(side_effect=Exception)):
-            self.zk.get_cluster()
         cluster = self.zk.get_cluster()
         self.assertEqual(cluster.last_lsn, 500)
 
@@ -295,3 +283,7 @@ class TestZooKeeper(unittest.TestCase):
 
     def test_set_history_value(self):
         self.zk.set_history_value('{}')
+
+    def test_watcher(self):
+        self.zk._watcher(WatchedEvent('', '', ''))
+        self.assertTrue(self.zk.watch(1, 1))