From a4a3d0b5f3ae07650086301a249c52900ab63e24 Mon Sep 17 00:00:00 2001
From: Michel Heily <michel.heily@granulate.io>
Date: Sun, 27 Aug 2023 11:47:40 +0300
Subject: [PATCH] Strictier type checking for auth

---
 glogger/sender.py             | 32 ++++++++++++++++++++++----------
 tests/glogger/test_handler.py |  6 +++---
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/glogger/sender.py b/glogger/sender.py
index 0575ffc1..6bf89b62 100644
--- a/glogger/sender.py
+++ b/glogger/sender.py
@@ -27,6 +27,15 @@ class SendBatch(NamedTuple):
     lost_logs_count: int
 
 
+class AuthToken(str):
+    pass
+
+
+class BasicAuthCredentials(NamedTuple):
+    username: str
+    password: str
+
+
 class Sender:
     # If Tuple[float, float], then the first value is connection-timeout and the second read-timeout.
     # See https://requests.readthedocs.io/en/latest/user/advanced/#timeouts
@@ -38,7 +47,7 @@ def __init__(
         server_address: str,
         *,
         auth_token: Optional[str] = None,
-        basic_auth: Optional[Tuple[str, str]] = None,
+        auth: Union[AuthToken, BasicAuthCredentials] = None,
         scheme: str = "https",
         send_interval: float = 30.0,
         send_threshold: float = 0.8,
@@ -50,8 +59,8 @@ def __init__(
         Create a new Sender and start flushing log messages in a background thread.
 
         :param application_name: Unique identifier requests coming from this handler.
-        :param auth_token: Token for authenticating requests to the server.
-        :param basic_auth: Basic auth credentials for authenticating requests to the server.
+        :param auth: The auth to use for this handler. One of AuthToken or BasicAuthCredentials.
+        :param auth_token: Deprecated - please use the auth param instead
         :param server_address: Address of server where to send messages.
         :param scheme: The scheme to use as string ('http' or 'https')
         :param send_interval: Seconds between sending batches.
@@ -76,13 +85,16 @@ def __init__(
         self.session = Session()
 
         # Set up auth
-        # basic-auth is preferred over X-Token
-        if basic_auth is not None:
-            self.session.auth = HTTPBasicAuth(*basic_auth)
-        elif auth_token is not None:
-            self.session.headers["X-Token"] = auth_token
-        else:
-            raise ValueError("Either auth_token or basic_auth must be provided")
+        # TODO: Remove auth_token in next version
+        if auth_token is not None:
+            import warnings
+
+            warnings.warn("auth_token is deprecated, please use the auth param", DeprecationWarning)
+
+        if isinstance(auth, BasicAuthCredentials):
+            self.session.auth = HTTPBasicAuth(*auth)
+        elif isinstance(auth, AuthToken):
+            self.session.headers["X-Token"] = str(auth)
 
         self.session.verify = verify
         self.messages_buffer: Optional[MessagesBuffer] = None
diff --git a/tests/glogger/test_handler.py b/tests/glogger/test_handler.py
index 38d1586f..c5f5a844 100644
--- a/tests/glogger/test_handler.py
+++ b/tests/glogger/test_handler.py
@@ -19,7 +19,7 @@
 
 from glogger.extra_adapter import ExtraAdapter
 from glogger.handler import BatchRequestsHandler
-from glogger.sender import SERVER_SEND_ERROR_MESSAGE, Sender
+from glogger.sender import SERVER_SEND_ERROR_MESSAGE, AuthToken, Sender
 
 
 class MockBatchRequestsHandler(BatchRequestsHandler):
@@ -33,7 +33,7 @@ def _send_once_to_server(self, data: bytes) -> None:
     def __init__(self, *args, max_total_length=100000, max_message_size=10000, overflow_drop_factor=0.25, **kwargs):
         super().__init__(
             self.MockSender(
-                "app", *args, auth_token="token", scheme="http", send_min_interval=0.2, max_send_tries=1, **kwargs
+                "app", *args, auth=AuthToken("token"), scheme="http", send_min_interval=0.2, max_send_tries=1, **kwargs
             ),
             max_total_length=max_total_length,
             max_message_size=max_message_size,
@@ -53,7 +53,7 @@ def __init__(self, *args, send_interval=0.2, **kwargs):
             self.HttpSender(
                 "app",
                 *args,
-                auth_token="token",
+                auth=AuthToken("token"),
                 scheme="http",
                 send_interval=send_interval,
                 send_min_interval=0.2,