From e39f97a2557dccab11a91824adee8afab2231d59 Mon Sep 17 00:00:00 2001
From: Antonio Andelic <antonio@clickhouse.com>
Date: Tue, 29 Nov 2022 09:13:23 +0000
Subject: [PATCH] Properly close connections and add retries

---
 .../test_keeper_zookeeper_converter/test.py   | 39 ++++++++++++++-----
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/tests/integration/test_keeper_zookeeper_converter/test.py b/tests/integration/test_keeper_zookeeper_converter/test.py
index af8d1ca4bf9..aa2e435ce36 100644
--- a/tests/integration/test_keeper_zookeeper_converter/test.py
+++ b/tests/integration/test_keeper_zookeeper_converter/test.py
@@ -2,14 +2,9 @@
 import pytest
 from helpers.cluster import ClickHouseCluster
 import helpers.keeper_utils as keeper_utils
-from kazoo.client import KazooClient, KazooState
-from kazoo.security import ACL, make_digest_acl, make_acl
-from kazoo.exceptions import (
-    AuthFailedError,
-    InvalidACLError,
-    NoAuthError,
-    KazooException,
-)
+from kazoo.client import KazooClient
+from kazoo.retry import KazooRetry
+from kazoo.security import make_acl
 import os
 import time
 
@@ -99,7 +94,9 @@ def get_fake_zk(timeout=60.0):
 
 def get_genuine_zk(timeout=60.0):
     _genuine_zk_instance = KazooClient(
-        hosts=cluster.get_instance_ip("node") + ":2181", timeout=timeout
+        hosts=cluster.get_instance_ip("node") + ":2181",
+        timeout=timeout,
+        connection_retry=KazooRetry(max_tries=20),
     )
     _genuine_zk_instance.start()
     return _genuine_zk_instance
@@ -225,6 +222,12 @@ def test_smoke(started_cluster, create_snapshots):
 
     compare_states(genuine_connection, fake_connection)
 
+    genuine_connection.stop()
+    genuine_connection.close()
+
+    fake_connection.stop()
+    fake_connection.close()
+
 
 def get_bytes(s):
     return s.encode()
@@ -309,6 +312,12 @@ def test_simple_crud_requests(started_cluster, create_snapshots):
     second_children = list(sorted(fake_connection.get_children("/test_sequential")))
     assert first_children == second_children, "Childrens are not equal on path " + path
 
+    genuine_connection.stop()
+    genuine_connection.close()
+
+    fake_connection.stop()
+    fake_connection.close()
+
 
 @pytest.mark.parametrize(("create_snapshots"), [True, False])
 def test_multi_and_failed_requests(started_cluster, create_snapshots):
@@ -379,6 +388,12 @@ def test_multi_and_failed_requests(started_cluster, create_snapshots):
     assert eph1 == eph2
     compare_stats(stat1, stat2, "/test_multitransactions", ignore_pzxid=True)
 
+    genuine_connection.stop()
+    genuine_connection.close()
+
+    fake_connection.stop()
+    fake_connection.close()
+
 
 @pytest.mark.parametrize(("create_snapshots"), [True, False])
 def test_acls(started_cluster, create_snapshots):
@@ -446,3 +461,9 @@ def test_acls(started_cluster, create_snapshots):
             "user2:lo/iTtNMP+gEZlpUNaCqLYO3i5U=",
             "user3:wr5Y0kEs9nFX3bKrTMKxrlcFeWo=",
         )
+
+    genuine_connection.stop()
+    genuine_connection.close()
+
+    fake_connection.stop()
+    fake_connection.close()