Make it possible to set a custom user agent

Author: leonid-shevtsov

mailtrap/client.py

@@ -1,7 +1,8 @@
 import warnings
-from typing import Optional
-from typing import Union
-from typing import cast
+
+import importlib.metadata
+
+from typing import Optional, Union, cast
 
 from pydantic import TypeAdapter
 
@@ -34,6 +35,10 @@ class MailtrapClient:
     DEFAULT_PORT = 443
     BULK_HOST = BULK_HOST
     SANDBOX_HOST = SANDBOX_HOST
+    DEFAULT_USER_AGENT = (
+        f"mailtrap-python/{importlib.metadata.version('mailtrap')} "
+        "(https://github.com/railsware/mailtrap-python)"
+    )
 
     def __init__(
         self,
@@ -44,6 +49,7 @@ def __init__(
         sandbox: bool = False,
         account_id: Optional[str] = None,
         inbox_id: Optional[str] = None,
+        user_agent: Optional[str] = None,
     ) -> None:
         self.token = token
         self.api_host = api_host
@@ -52,6 +58,9 @@ def __init__(
         self.sandbox = sandbox
         self.account_id = account_id
         self.inbox_id = inbox_id
+        self._user_agent = (
+            user_agent if user_agent is not None else self.DEFAULT_USER_AGENT
+        )
 
         self._validate_itself()
 
@@ -147,9 +156,7 @@ def headers(self) -> dict[str, str]:
         return {
             "Authorization": f"Bearer {self.token}",
             "Content-Type": "application/json",
-            "User-Agent": (
-                "mailtrap-python (https://github.com/railsware/mailtrap-python)"
-            ),
+            "User-Agent": self._user_agent,
         }
 
     @property

tests/unit/test_client.py

@@ -86,7 +86,11 @@ def test_headers_should_return_appropriate_dict(self) -> None:
         assert client.headers == {
             "Authorization": "Bearer fake_token",
             "Content-Type": "application/json",
-            "User-Agent": (
-                "mailtrap-python (https://github.com/railsware/mailtrap-python)"
-            ),
+            "User-Agent": mt.MailtrapClient.DEFAULT_USER_AGENT,
         }
+
+    def test_headers_should_use_custom_user_agent_when_provided(self) -> None:
+        custom_ua = "MyApp/1.0 (custom)"
+        client = self.get_client(user_agent=custom_ua)
+
+        assert client.headers["User-Agent"] == custom_ua