From 60d7493821e90544ea9916f9a5b15637a49d3377 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-S=C3=A9bastien=20Sevestre?=
 <jean-sebastien@lumapps.com>
Date: Thu, 16 Feb 2023 17:57:55 +0100
Subject: [PATCH] feat(api): keep extra headers in get_new_client_as

---
 lumapps/api/base_client.py  |  8 ++++++--
 tests/legacy/test_client.py | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/lumapps/api/base_client.py b/lumapps/api/base_client.py
index 576cfe65..ceea03ed 100644
--- a/lumapps/api/base_client.py
+++ b/lumapps/api/base_client.py
@@ -74,8 +74,9 @@ def __init__(
         self._endpoints = None
         self._client = None
         self._headers: dict = {"x-lumapps-analytics": "off"}
-        if extra_http_headers:
-            self._headers.update(extra_http_headers)
+        self._extra_http_headers = extra_http_headers
+        if self._extra_http_headers:
+            self._headers.update(self._extra_http_headers)
         if api_info is None:
             api_info = {}
         api_info.setdefault("name", LUMAPPS_NAME)
@@ -231,6 +232,7 @@ def get_new_client_as_using_dwd(self, user_email: str) -> "BaseClient":
             no_verify=self.no_verify,
             proxy_info=self.proxy_info,
             prune=self.prune,
+            extra_http_headers=self._extra_http_headers,
         )
 
     def get_new_client_as(
@@ -252,6 +254,7 @@ def get_new_client_as(
             no_verify=self.no_verify,
             proxy_info=self.proxy_info,
             prune=self.prune,
+            extra_http_headers=self._extra_http_headers,
         )
         token_infos: Any = client.get_call(
             "user/getToken", customerId=customer_id, email=user_email
@@ -264,6 +267,7 @@ def get_new_client_as(
             no_verify=self.no_verify,
             proxy_info=self.proxy_info,
             prune=self.prune,
+            extra_http_headers=self._extra_http_headers,
         )
 
     @property
diff --git a/tests/legacy/test_client.py b/tests/legacy/test_client.py
index 3139426b..f2e021ad 100644
--- a/tests/legacy/test_client.py
+++ b/tests/legacy/test_client.py
@@ -43,3 +43,17 @@ def test_custom_headers():
     headers = {"my-header": "on"}
     cli = LumAppsClient("a", "b", token="FAKE", extra_http_headers=headers)
     assert "my-header" in cli.client.headers
+
+
+def test_custom_headers_for_new_client(mocker):
+    def dummy_get_call(*args, **kwargs):
+        return {"accessToken": "1"}
+    mocker.patch(
+        "lumapps.api.base_client.BaseClient.get_call",
+        dummy_get_call,
+    )
+
+    headers = {"my-header": "on"}
+    cli = LumAppsClient("a", "b", token="FAKE", extra_http_headers=headers)
+    new_as = cli.get_new_client_as("user_email")
+    assert "my-header" in new_as.client.headers