From 0e46cf86b772e1513d837d6019181a6d291b7219 Mon Sep 17 00:00:00 2001
From: Smita Kulkarni <Smita.Kulkarni@clickhouse.com>
Date: Mon, 24 Jul 2023 08:52:19 +0200
Subject: [PATCH] Added try-except to check cases when second backup/restore is
 picked up first

---
 .../test_disallow_concurrency.py              | 69 +++++++++++++++----
 1 file changed, 57 insertions(+), 12 deletions(-)

diff --git a/tests/integration/test_backup_restore_on_cluster/test_disallow_concurrency.py b/tests/integration/test_backup_restore_on_cluster/test_disallow_concurrency.py
index d0ce2e03016..a863a6e2047 100644
--- a/tests/integration/test_backup_restore_on_cluster/test_disallow_concurrency.py
+++ b/tests/integration/test_backup_restore_on_cluster/test_disallow_concurrency.py
@@ -133,9 +133,21 @@ def test_concurrent_backups_on_same_node():
     )
     assert status in ["CREATING_BACKUP", "BACKUP_CREATED"]
 
-    error = nodes[0].query_and_get_error(
-        f"BACKUP TABLE tbl ON CLUSTER 'cluster' TO {backup_name}"
-    )
+    try:
+        error = nodes[0].query_and_get_error(
+            f"BACKUP TABLE tbl ON CLUSTER 'cluster' TO {backup_name}"
+        )
+    except Exception as e:
+        status = (
+            nodes[0]
+            .query(f"SELECT status FROM system.backups WHERE id == '{id}'")
+            .rstrip("\n")
+        )
+        # It is possible that the second backup was picked up first, and then the async backup
+        if status == "CREATING_BACKUP" or status == "BACKUP_FAILED":
+            return
+        else:
+            raise e
     expected_errors = [
         "Concurrent backups not supported",
         f"Backup {backup_name} already exists",
@@ -179,9 +191,20 @@ def test_concurrent_backups_on_different_nodes():
     )
     assert status in ["CREATING_BACKUP", "BACKUP_CREATED"]
 
-    error = nodes[0].query_and_get_error(
-        f"BACKUP TABLE tbl ON CLUSTER 'cluster' TO {backup_name}"
-    )
+    try:
+        error = nodes[0].query_and_get_error(
+            f"BACKUP TABLE tbl ON CLUSTER 'cluster' TO {backup_name}"
+        )
+    except Exception as e:
+        status = (
+            nodes[1]
+            .query(f"SELECT status FROM system.backups WHERE id == '{id}'")
+            .rstrip("\n")
+        )
+        if status == "CREATING_BACKUP" or status == "BACKUP_FAILED":
+            return
+        else:
+            raise e
     expected_errors = [
         "Concurrent backups not supported",
         f"Backup {backup_name} already exists",
@@ -224,9 +247,20 @@ def test_concurrent_restores_on_same_node():
     )
     assert status in ["RESTORING", "RESTORED"]
 
-    error = nodes[0].query_and_get_error(
-        f"RESTORE TABLE tbl ON CLUSTER 'cluster' FROM {backup_name}"
-    )
+    try:
+        error = nodes[0].query_and_get_error(
+            f"RESTORE TABLE tbl ON CLUSTER 'cluster' FROM {backup_name}"
+        )
+    except Exception as e:
+        status = (
+            nodes[0]
+            .query(f"SELECT status FROM system.backups WHERE id == '{id}'")
+            .rstrip("\n")
+        )
+        if status == "RESTORING" or status == "RESTORE_FAILED":
+            return
+        else:
+            raise e
     expected_errors = [
         "Concurrent restores not supported",
         "Cannot restore the table default.tbl because it already contains some data",
@@ -269,9 +303,20 @@ def test_concurrent_restores_on_different_node():
     )
     assert status in ["RESTORING", "RESTORED"]
 
-    error = nodes[1].query_and_get_error(
-        f"RESTORE TABLE tbl ON CLUSTER 'cluster' FROM {backup_name}"
-    )
+    try:
+        error = nodes[1].query_and_get_error(
+            f"RESTORE TABLE tbl ON CLUSTER 'cluster' FROM {backup_name}"
+        )
+    except Exception as e:
+        status = (
+            nodes[0]
+            .query(f"SELECT status FROM system.backups WHERE id == '{id}'")
+            .rstrip("\n")
+        )
+        if status == "RESTORING" or status == "RESTORE_FAILED":
+            return
+        else:
+            raise e
     expected_errors = [
         "Concurrent restores not supported",
         "Cannot restore the table default.tbl because it already contains some data",