From 7cadfeac29c68dc9ac749cd48d115da02ae48f62 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A1nos=20Benjamin=20Antal?=
 <benjamin.antal@clickhouse.com>
Date: Mon, 26 Jun 2023 14:55:03 +0000
Subject: [PATCH] Add retries to flaky tests

---
 tests/integration/test_tlsv1_3/test.py | 67 ++++++++++++++++++--------
 1 file changed, 47 insertions(+), 20 deletions(-)

diff --git a/tests/integration/test_tlsv1_3/test.py b/tests/integration/test_tlsv1_3/test.py
index e25216c67df..a29dcb5fc1c 100644
--- a/tests/integration/test_tlsv1_3/test.py
+++ b/tests/integration/test_tlsv1_3/test.py
@@ -4,6 +4,7 @@ from helpers.ssl_context import WrapSSLContextWithSNI
 import urllib.request, urllib.parse
 import ssl
 import os.path
+import logging
 
 
 # The test cluster is configured with certificate for that host name, see 'server-ext.cnf'.
@@ -11,6 +12,7 @@ import os.path
 SSL_HOST = "integration-tests.clickhouse.com"
 HTTPS_PORT = 8443
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
+MAX_RETRY = 5
 
 cluster = ClickHouseCluster(__file__)
 instance = cluster.add_instance(
@@ -88,10 +90,18 @@ def test_https_wrong_cert():
         execute_query_https("SELECT currentUser()", user="john", cert_name="client2")
     assert "HTTP Error 403" in str(err.value)
 
+    count = 0
     # Wrong certificate: self-signed certificate.
-    with pytest.raises(Exception) as err:
-        execute_query_https("SELECT currentUser()", user="john", cert_name="wrong")
-    assert "unknown ca" in str(err.value)
+    while count <= MAX_RETRY:
+        with pytest.raises(Exception) as err:
+            execute_query_https("SELECT currentUser()", user="john", cert_name="wrong")
+        err_str = str(err.value)
+        if count < MAX_RETRY and "Broken pipe" in err_str:
+            count = count + 1
+            logging.warning(f"Failed attempt with wrong cert, err: {err_str}")
+            continue
+        assert "unknown ca" in err_str
+        break
 
     # No certificate.
     with pytest.raises(Exception) as err:
@@ -181,24 +191,41 @@ def test_https_non_ssl_auth():
         == "jane\n"
     )
 
+    count = 0
     # However if we send a certificate it must not be wrong.
-    with pytest.raises(Exception) as err:
-        execute_query_https(
-            "SELECT currentUser()",
-            user="peter",
-            enable_ssl_auth=False,
-            cert_name="wrong",
-        )
-    assert "unknown ca" in str(err.value)
-    with pytest.raises(Exception) as err:
-        execute_query_https(
-            "SELECT currentUser()",
-            user="jane",
-            enable_ssl_auth=False,
-            password="qwe123",
-            cert_name="wrong",
-        )
-    assert "unknown ca" in str(err.value)
+    while count <= MAX_RETRY:
+        with pytest.raises(Exception) as err:
+            execute_query_https(
+                "SELECT currentUser()",
+                user="peter",
+                enable_ssl_auth=False,
+                cert_name="wrong",
+            )
+        err_str = str(err.value)
+        if count < MAX_RETRY and "Broken pipe" in err_str:
+            count = count + 1
+            logging.warning(f"Failed attempt with wrong cert, user: peter, err: {err_str}")
+            continue
+        assert "unknown ca" in err_str
+        break
+
+    count = 0
+    while count <= MAX_RETRY:
+        with pytest.raises(Exception) as err:
+            execute_query_https(
+                "SELECT currentUser()",
+                user="jane",
+                enable_ssl_auth=False,
+                password="qwe123",
+                cert_name="wrong",
+            )
+        err_str = str(err.value)
+        if count < MAX_RETRY and "Broken pipe" in err_str:
+            count = count + 1
+            logging.warning(f"Failed attempt with wrong cert, user: jane, err: {err_str}")
+            continue
+        assert "unknown ca" in err_str
+        break
 
 
 def test_create_user():