Merge pull request #21779 from MaterializeInc/roshan/request-headers

[Cloudtest] Allow setting additional headers in cloudtest requests
MaterializeInc · Sep 18, 2023 · 556dcce · 556dcce
2 parents adfbfb7 + 3e8259e
commit 556dcce
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/misc/python/materialize/cloudtest/util/controller.py b/misc/python/materialize/cloudtest/util/controller.py
@@ -58,9 +58,17 @@ def configured_base_url(self) -> str:
         return self.endpoint.base_url
 
     def requests(
-        self, auth: AuthConfig | None, client_cert: tuple[str, str] | None = None
+        self,
+        auth: AuthConfig | None,
+        client_cert: tuple[str, str] | None = None,
+        additional_headers: dict[str, str] | None = None,
     ) -> WebRequests:
-        return WebRequests(auth, self.configured_base_url(), client_cert)
+        return WebRequests(
+            auth,
+            self.configured_base_url(),
+            client_cert=client_cert,
+            additional_headers=additional_headers,
+        )
 
 
 def wait_for_connectable(

diff --git a/misc/python/materialize/cloudtest/util/web_request.py b/misc/python/materialize/cloudtest/util/web_request.py
@@ -38,10 +38,12 @@ def __init__(
         auth: AuthConfig | None,
         base_url: str,
         client_cert: tuple[str, str] | None = None,
+        additional_headers: dict[str, str] | None = None,
     ):
         self.auth = auth
         self.base_url = base_url
         self.client_cert = client_cert
+        self.additional_headers = additional_headers
 
     def get(
         self,
@@ -176,7 +178,7 @@ def try_delete() -> requests.Response:
         return response
 
     def _create_headers(self, auth: AuthConfig | None) -> dict[str, Any]:
-        headers = {}
+        headers = self.additional_headers.copy() if self.additional_headers else {}
         if auth:
             headers["Authorization"] = f"Bearer {auth.token}"