From ce29d1588b1c0acaca2ea115787a3cdaa52d65b6 Mon Sep 17 00:00:00 2001
From: John Spray <john@neon.tech>
Date: Fri, 13 Oct 2023 08:34:50 +0100
Subject: [PATCH] tests: add helpers for location_conf

---
 test_runner/fixtures/neon_fixtures.py   | 24 +++++++++++++++++++++++-
 test_runner/fixtures/pageserver/http.py |  8 ++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index bdb6ae2370..f3cb1aa370 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -28,6 +28,7 @@ import jwt
 import psycopg2
 import pytest
 import requests
+import toml
 from _pytest.config import Config
 from _pytest.config.argparsing import Parser
 from _pytest.fixtures import FixtureRequest
@@ -1859,6 +1860,27 @@ class NeonPageserver(PgProtocol):
         client = self.http_client()
         return client.tenant_attach(tenant_id, config, config_null, generation=generation)
 
+    def tenant_location_configure(self, tenant_id: TenantId, config: dict[str, Any]):
+        # This API is only for use when generations are enabled
+        assert self.env.attachment_service is not None
+
+        if config["mode"].startswith("Attached") and "generation" not in config:
+            config["generation"] = self.env.attachment_service.attach_hook(tenant_id, self.id)
+
+        client = self.http_client()
+        return client.tenant_location_conf(tenant_id, config)
+
+    def read_tenant_location_conf(self, tenant_id: TenantId) -> dict[str, Any]:
+        path = self.tenant_dir(tenant_id) / "config-v1"
+        log.info(f"Reading location conf from {path}")
+        bytes = open(path, "r").read()
+        try:
+            decoded: dict[str, Any] = toml.loads(bytes)
+            return decoded
+        except:
+            log.error(f"Failed to decode LocationConf, raw content ({len(bytes)} bytes): {bytes}")
+            raise
+
 
 def append_pageserver_param_overrides(
     params_to_update: List[str],
@@ -3079,7 +3101,7 @@ def pytest_addoption(parser: Parser):
 
 
 SMALL_DB_FILE_NAME_REGEX: re.Pattern = re.compile(  # type: ignore[type-arg]
-    r"config|metadata|.+\.(?:toml|pid|json|sql)"
+    r"config|config-v1|heatmap-v1|metadata|.+\.(?:toml|pid|json|sql)"
 )
 
 
diff --git a/test_runner/fixtures/pageserver/http.py b/test_runner/fixtures/pageserver/http.py
index aff7959aa7..ae11859065 100644
--- a/test_runner/fixtures/pageserver/http.py
+++ b/test_runner/fixtures/pageserver/http.py
@@ -251,6 +251,14 @@ class PageserverHttpClient(requests.Session):
         res = self.post(f"http://localhost:{self.port}/v1/tenant/{tenant_id}/detach", params=params)
         self.verbose_error(res)
 
+    def tenant_location_conf(self, tenant_id: TenantId, location_conf=dict[str, Any]):
+        body = location_conf.copy()
+        body["tenant_id"] = str(tenant_id)
+        res = self.put(
+            f"http://localhost:{self.port}/v1/tenant/{tenant_id}/location_config", json=body
+        )
+        self.verbose_error(res)
+
     def tenant_delete(self, tenant_id: TenantId):
         res = self.delete(f"http://localhost:{self.port}/v1/tenant/{tenant_id}")
         self.verbose_error(res)