From c6f9b8f3183c42219d5d03ac2324891d6472a272 Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Fri, 26 May 2023 17:37:42 +0200
Subject: [PATCH] test_broken_timelines: wait for tenants to load

Without this, we rely on the basebackup request to wait for the tenant to load.

It works, but, would be nice to rule it out, no?
---
 test_runner/fixtures/pageserver/utils.py    | 14 ++++++++++++++
 test_runner/regress/test_broken_timeline.py |  6 ++++++
 2 files changed, 20 insertions(+)

diff --git a/test_runner/fixtures/pageserver/utils.py b/test_runner/fixtures/pageserver/utils.py
index c558387413..2818109922 100644
--- a/test_runner/fixtures/pageserver/utils.py
+++ b/test_runner/fixtures/pageserver/utils.py
@@ -4,6 +4,7 @@ from typing import Optional
 from fixtures.log_helper import log
 from fixtures.pageserver.http import PageserverHttpClient
 from fixtures.types import Lsn, TenantId, TimelineId
+from fixtures.utils import wait_until
 
 
 def assert_tenant_state(
@@ -107,6 +108,19 @@ def wait_until_tenant_active(
     )
 
 
+def assert_all_tenants_to_exit_loading_state(pageserver_http: PageserverHttpClient):
+    states = [t["state"]["slug"] != "Loading" for t in pageserver_http.tenant_list()]
+    assert all(states)
+
+
+def wait_for_all_tenants_to_exit_loading_state(
+    pageserver_http: PageserverHttpClient, iterations: int = 25, interval: float = 0.2
+):
+    wait_until(
+        iterations, interval, lambda: assert_all_tenants_to_exit_loading_state(pageserver_http)
+    )
+
+
 def last_record_lsn(
     pageserver_http_client: PageserverHttpClient, tenant: TenantId, timeline: TimelineId
 ) -> Lsn:
diff --git a/test_runner/regress/test_broken_timeline.py b/test_runner/regress/test_broken_timeline.py
index 1ce9fd006f..3618a4e400 100644
--- a/test_runner/regress/test_broken_timeline.py
+++ b/test_runner/regress/test_broken_timeline.py
@@ -5,6 +5,9 @@ from typing import List, Tuple
 import pytest
 from fixtures.log_helper import log
 from fixtures.neon_fixtures import Endpoint, NeonEnv, NeonEnvBuilder
+from fixtures.pageserver.utils import (
+    wait_for_all_tenants_to_exit_loading_state,
+)
 from fixtures.types import TenantId, TimelineId
 
 
@@ -12,6 +15,7 @@ from fixtures.types import TenantId, TimelineId
 # running.
 def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
     env = neon_env_builder.init_start()
+    client = env.pageserver.http_client()
 
     env.pageserver.allowed_errors.extend(
         [
@@ -72,6 +76,8 @@ def test_broken_timeline(neon_env_builder: NeonEnvBuilder):
 
     env.pageserver.start()
 
+    wait_for_all_tenants_to_exit_loading_state(client)
+
     # Tenant 0 should still work
     pg0.start()
     assert pg0.safe_psql("SELECT COUNT(*) FROM t")[0][0] == 100