From 74626b94a8020232d95b2655ca6256846924abd6 Mon Sep 17 00:00:00 2001 From: Alexey Masterov Date: Thu, 17 Jul 2025 18:48:44 +0200 Subject: [PATCH] An attempt 2 --- test_runner/random_ops/test_random_ops.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/test_runner/random_ops/test_random_ops.py b/test_runner/random_ops/test_random_ops.py index 59521dc99a..ca5d806e06 100644 --- a/test_runner/random_ops/test_random_ops.py +++ b/test_runner/random_ops/test_random_ops.py @@ -504,17 +504,24 @@ class NeonProject: Creates a new Neon branch for the current project, then restores the snapshot with the given id """ - target_branch = self.create_branch(is_reset=True) + target_branch = self.create_branch() if not target_branch: return None - new_branch_id = self.neon_api.restore_snapshot( + new_branch_def = self.neon_api.restore_snapshot( self.id, snapshot_id, target_branch.id, self.generate_branch_name(), - )["branch"]["id"] + ) self.wait() - new_branch = NeonBranch(self, self.neon_api.get_branch_details(self.id, new_branch_id)) + NeonBranch( + self, + self.neon_api.get_branch_details(self.id, new_branch_def["branch"]["parent_id"]), + is_reset=True, + ) + new_branch = NeonBranch( + self, self.neon_api.get_branch_details(self.id, new_branch_def["branch"]["id"]) + ) if new_branch.connection_parameters is None: raise RuntimeError(f"Neon branch {new_branch.id} does not have connection parameters") with psycopg2.connect( @@ -531,7 +538,9 @@ class NeonProject: snapshot_name = row[0] assert snapshot_name == self.snapshots[snapshot_id].name self.wait() - return target_branch + target_branch.start_benchmark() + new_branch.start_benchmark() + return new_branch @pytest.fixture()