From 018f95115ae2f748e780f32a9d6c241c68849360 Mon Sep 17 00:00:00 2001
From: Alexey Masterov <alexey.masterov@databricks.com>
Date: Wed, 16 Jul 2025 13:21:32 +0200
Subject: [PATCH] Retry on 423 error "snapshot is in transition"

---
 test_runner/fixtures/neon_api.py          | 1 +
 test_runner/random_ops/test_random_ops.py | 6 +++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/test_runner/fixtures/neon_api.py b/test_runner/fixtures/neon_api.py
index c04d4600de..122fe5186d 100644
--- a/test_runner/fixtures/neon_api.py
+++ b/test_runner/fixtures/neon_api.py
@@ -66,6 +66,7 @@ class NeonAPI:
                 elif resp.status_code == 423 and resp.json()["message"] in {
                     "endpoint is in some transitive state, could not suspend",
                     "project already has running conflicting operations, scheduling of new ones is prohibited",
+                    "snapshot is in transition"
                 }:
                     retry = True
                     self.retries4xx += 1
diff --git a/test_runner/random_ops/test_random_ops.py b/test_runner/random_ops/test_random_ops.py
index c29b2e97c7..62e39ce6bb 100644
--- a/test_runner/random_ops/test_random_ops.py
+++ b/test_runner/random_ops/test_random_ops.py
@@ -463,7 +463,7 @@ class NeonProject:
         Two optional arguments: lsn and timestamp are mutually exclusive
         they instruct to create a snapshot with the specific lns or timestamp
         """
-        return NeonSnapshot(
+        snapshot = NeonSnapshot(
             self,
             self.neon_api.create_snapshot(
                 self.id,
@@ -473,6 +473,8 @@ class NeonProject:
                 self.gen_snapshot_name(),
             ),
         )
+        self.wait()
+        return snapshot
 
     def delete_snapshot(self, snapshot_id: str) -> None:
         """
@@ -480,6 +482,7 @@ class NeonProject:
         """
         self.neon_api.delete_snapshot(self.id, snapshot_id)
         self.snapshots.pop(snapshot_id)
+        self.wait()
 
     def restore_snapshot(self, snapshot_id: str) -> NeonBranch | None:
         """
@@ -496,6 +499,7 @@ class NeonProject:
             self.generate_branch_name(),
         )
         target_branch.start_benchmark()
+        self.wait()
         return target_branch