Fix review: Use AuthType enum values

Signed-off-by: Viet Nguyen Duc <[email protected]>

Author: VietND96

py/selenium/webdriver/remote/client_config.py

@@ -17,6 +17,7 @@
 import base64
 import os
 import socket
+from enum import Enum
 from typing import Optional
 from urllib import parse
 
@@ -26,6 +27,12 @@
 from selenium.webdriver.common.proxy import ProxyType
 
 
+class AuthType(Enum):
+    BASIC = "Basic"
+    BEARER = "Bearer"
+    X_API_KEY = "X-API-Key"
+
+
 class ClientConfig:
     def __init__(
         self,
@@ -38,7 +45,7 @@ def __init__(
         ca_certs: Optional[str] = None,
         username: Optional[str] = None,
         password: Optional[str] = None,
-        auth_type: Optional[str] = "Basic",
+        auth_type: Optional[AuthType] = AuthType.BASIC,
         token: Optional[str] = None,
         user_agent: Optional[str] = None,
         extra_headers: Optional[dict] = None,
@@ -202,16 +209,16 @@ def password(self, value: str) -> None:
         self._password = value
 
     @property
-    def auth_type(self) -> str:
+    def auth_type(self) -> AuthType:
         """Returns the type of authentication to the remote server."""
         return self._auth_type
 
     @auth_type.setter
-    def auth_type(self, value: str) -> None:
+    def auth_type(self, value: AuthType) -> None:
         """Sets the type of authentication to the remote server if it is not
         using basic with username and password.
 
-        Support values: Bearer, X-API-Key. For others, please use `extra_headers` instead
+        :Args: value - AuthType enum value. For others, please use `extra_headers` instead
         """
         self._auth_type = value
 
@@ -273,13 +280,12 @@ def get_proxy_url(self) -> Optional[str]:
 
     def get_auth_header(self) -> Optional[dict]:
         """Returns the authorization to add to the request headers."""
-        auth_type = self.auth_type.lower()
-        if auth_type == "basic" and self.username and self.password:
+        if self.auth_type is AuthType.BASIC and self.username and self.password:
             credentials = f"{self.username}:{self.password}"
             encoded_credentials = base64.b64encode(credentials.encode("utf-8")).decode("utf-8")
-            return {"Authorization": f"Basic {encoded_credentials}"}
-        if auth_type == "bearer" and self.token:
-            return {"Authorization": f"Bearer {self.token}"}
-        if auth_type == "x-api-key" and self.token:
-            return {"X-API-Key": f"{self.token}"}
+            return {"Authorization": f"{AuthType.BASIC.value} {encoded_credentials}"}
+        if self.auth_type is AuthType.BEARER and self.token:
+            return {"Authorization": f"{AuthType.BEARER.value} {self.token}"}
+        if self.auth_type is AuthType.X_API_KEY and self.token:
+            return {f"{AuthType.X_API_KEY.value}": f"{self.token}"}
         return None

py/test/unit/selenium/webdriver/remote/remote_connection_tests.py

@@ -27,6 +27,7 @@
 from selenium import __version__
 from selenium.webdriver import Proxy
 from selenium.webdriver.common.proxy import ProxyType
+from selenium.webdriver.remote.client_config import AuthType
 from selenium.webdriver.remote.remote_connection import ClientConfig
 from selenium.webdriver.remote.remote_connection import RemoteConnection
 
@@ -93,7 +94,7 @@ def test_get_proxy_url_http(mock_proxy_settings):
 
 def test_get_auth_header_if_client_config_pass_basic_auth():
     custom_config = ClientConfig(
-        remote_server_addr="http://remote", keep_alive=True, username="user", password="pass", auth_type="Basic"
+        remote_server_addr="http://remote", keep_alive=True, username="user", password="pass", auth_type=AuthType.BASIC
     )
     remote_connection = RemoteConnection(custom_config.remote_server_addr, client_config=custom_config)
     headers = remote_connection._client_config.get_auth_header()
@@ -102,7 +103,7 @@ def test_get_auth_header_if_client_config_pass_basic_auth():
 
 def test_get_auth_header_if_client_config_pass_bearer_token():
     custom_config = ClientConfig(
-        remote_server_addr="http://remote", keep_alive=True, auth_type="Bearer", token="dXNlcjpwYXNz"
+        remote_server_addr="http://remote", keep_alive=True, auth_type=AuthType.BEARER, token="dXNlcjpwYXNz"
     )
     remote_connection = RemoteConnection(custom_config.remote_server_addr, client_config=custom_config)
     headers = remote_connection._client_config.get_auth_header()
@@ -111,7 +112,7 @@ def test_get_auth_header_if_client_config_pass_bearer_token():
 
 def test_get_auth_header_if_client_config_pass_x_api_key():
     custom_config = ClientConfig(
-        remote_server_addr="http://remote", keep_alive=True, auth_type="X-API-Key", token="abcdefgh123456789"
+        remote_server_addr="http://remote", keep_alive=True, auth_type=AuthType.X_API_KEY, token="abcdefgh123456789"
     )
     remote_connection = RemoteConnection(custom_config.remote_server_addr, client_config=custom_config)
     headers = remote_connection._client_config.get_auth_header()