From 8b0ee2abf962827cb434a9c6700fcd7f48c81f79 Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Tue, 17 Oct 2023 19:56:29 +0100
Subject: [PATCH] tests: add a quiesce helper for waiting for attachment

---
 test_runner/fixtures/neon_fixtures.py         | 36 +++++++++++++++++++
 .../regress/test_threshold_based_eviction.py  |  3 +-
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index a9fe5e376e..a48c05a4d0 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -60,6 +60,7 @@ from fixtures.utils import (
     allure_attach_from_dir,
     get_self_dir,
     subprocess_capture,
+    wait_until,
 )
 
 """
@@ -1680,6 +1681,41 @@ class NeonPageserver(PgProtocol):
             self.running = False
         return self
 
+    def restart(self, immediate: bool = False):
+        """
+        High level wrapper for restart: restarts the process, and waits for
+        tenant state to stabilize.
+        """
+        self.stop(immediate=immediate)
+        self.start()
+        self.quiesce_tenants()
+
+    def quiesce_tenants(self):
+        """
+        Wait for all tenants to enter a stable state (Active or Broken)
+
+        Call this after restarting the pageserver, or after attaching a tenant,
+        to ensure that it is ready for use.
+        """
+
+        stable_states = {"Active", "Broken"}
+
+        client = self.http_client()
+
+        def complete():
+            log.info("Checking tenants...")
+            tenants = client.tenant_list()
+            tenants = client.tenant_list()
+            log.info(f"Tenant list: {tenants}...")
+            any_unstable = any((t["state"]["slug"] not in stable_states) for t in tenants)
+            if any_unstable:
+                for t in tenants:
+                    log.info(f"Waiting for tenant {t['id']} in state {t['state']['slug']}")
+            log.info(f"any_unstable={any_unstable}")
+            assert not any_unstable
+
+        wait_until(20, 0.5, complete)
+
     def __enter__(self) -> "NeonPageserver":
         return self
 
diff --git a/test_runner/regress/test_threshold_based_eviction.py b/test_runner/regress/test_threshold_based_eviction.py
index 12866accc7..27d5cce5f2 100644
--- a/test_runner/regress/test_threshold_based_eviction.py
+++ b/test_runner/regress/test_threshold_based_eviction.py
@@ -70,8 +70,7 @@ def test_threshold_based_eviction(
     }
 
     # restart because changing tenant config is not instant
-    env.pageserver.stop()
-    env.pageserver.start()
+    env.pageserver.restart()
 
     assert ps_http.tenant_config(tenant_id).effective_config["eviction_policy"] == {
         "kind": "LayerAccessThreshold",