From bb1a50966191cb2f855c27ea6322394dd0d60d36 Mon Sep 17 00:00:00 2001
From: Eamonn Rea <eamonnrea@gmail.com>
Date: Thu, 29 Aug 2024 17:33:21 +0100
Subject: [PATCH] test_util: Add test for `build_headers_with_authorization`
 (#443)

* Add doctstring and typing to `build_headers_with_authorization`

* util: Use copy of request_headers for build_headers_with_authorization

* test_util: Add test for build_headers_with_authorization
---
 pupgui2/pupgui2.py |  2 +-
 pupgui2/util.py    | 21 +++++++++++++++------
 tests/test_util.py | 36 ++++++++++++++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/pupgui2/pupgui2.py b/pupgui2/pupgui2.py
index 3793324d..28eda2cc 100644
--- a/pupgui2/pupgui2.py
+++ b/pupgui2/pupgui2.py
@@ -82,7 +82,7 @@ class MainWindow(QObject):
     def __init__(self):
         super(MainWindow, self).__init__()
 
-        self.web_access_tokens = {
+        self.web_access_tokens: dict[str, str] = {
             'github': os.getenv('PUPGUI_GHA_TOKEN') or config_github_access_token(),
             'gitlab': os.getenv('PUPGUI_GLA_TOKEN') or config_gitlab_access_token(),
         }
diff --git a/pupgui2/util.py b/pupgui2/util.py
index 33ad5327..48881a87 100644
--- a/pupgui2/util.py
+++ b/pupgui2/util.py
@@ -666,19 +666,28 @@ def fetch_project_release_data(release_url: str, release_format: str, rs: reques
     return values
 
 
-def build_headers_with_authorization(request_headers: dict, authorization_tokens: dict, token_type: str):
+def build_headers_with_authorization(request_headers: dict[str, Any], authorization_tokens: dict[str, str], token_type: str) -> dict[str, Any]:
 
-    request_headers['Authorization'] = ''  # Reset old authentication
+    """
+    Generate an updated `request_headers` dict with the `Authorization` header containing the key for GitHub or GitLab, based on `token_type`
+    and removing any existing Authorization.
+
+    Return Type: dict[str, Any]
+    """
+
+    updated_headers: dict[str, Any] = request_headers.copy()
+
+    updated_headers['Authorization'] = ''  # Reset old authentication
     token: str = authorization_tokens.get(token_type, '')
     if not token:        
-        return request_headers
+        return updated_headers
 
     if token_type == 'github':
-        request_headers['Authorization'] = f'token {token}'
+        updated_headers['Authorization'] = f'token {token}'
     elif token_type == 'gitlab':
-        request_headers['Authorization'] = f'Bearer {token}'
+        updated_headers['Authorization'] = f'Bearer {token}'
 
-    return request_headers
+    return updated_headers
 
 def compat_tool_available(compat_tool: str, ctobjs: List[dict]) -> bool:
     """ Return whether a compat tool is available for a given launcher """
diff --git a/tests/test_util.py b/tests/test_util.py
index 6e879f88..1fffae12 100644
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -4,6 +4,42 @@
 from pupgui2.datastructures import SteamApp, LutrisGame, HeroicGame, Launcher
 
 
+def test_build_headers_with_authorization() -> None:
+
+    """
+    Test whether the expected Authorization Tokens get inserted into the returned headers dict,
+    that existing Authorization is replaced properly, and that all other existing headers are preserved.
+    """
+
+    user_agent = 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.103 Safari/537.36'
+
+    # Simulate existing headers with old Authentiation to be replaced, and a User-Agent that should remain untouched
+    request_headers: dict[str, Any] = {
+        'Authorization': 'ABC123',
+        'User-Agent': user_agent
+    }
+
+    # Simulate auth tokens that would normally come from the environment or config file
+    authorization_tokens: dict[str, str] = {
+        'github': 'gha_abc123daf456',
+        'gitlab': 'glpat-zyx987wvu654',
+    }
+
+    github_token_call: dict[str, Any] = build_headers_with_authorization(request_headers, authorization_tokens, 'github')
+    gitlab_token_call: dict[str, Any] = build_headers_with_authorization(request_headers, authorization_tokens, 'gitlab')
+
+    unknown_token_call: dict[str, Any] = build_headers_with_authorization(request_headers, authorization_tokens, '')
+    call_with_no_tokens: dict[str, Any] = build_headers_with_authorization(request_headers, {}, 'github')
+
+    assert github_token_call.get('Authorization', '') == f'token {authorization_tokens["github"]}'
+    assert gitlab_token_call.get('Authorization', '') == f'Bearer {authorization_tokens["gitlab"]}'
+
+    assert unknown_token_call.get('Authorization', '') == ''
+    assert call_with_no_tokens.get('Authorization', '') == ''
+
+    assert github_token_call.get('User-Agent', '') == user_agent
+
+
 def test_get_dict_key_from_value() -> None:
 
     """