Provide __eq__ implementations for subclasses of Authentication

d-cole · Dec 2, 2021 · 61de1e1 · 61de1e1
1 parent 367d665
commit 61de1e1
Showing 1 changed file with 28 additions and 0 deletions.
diff --git a/trino/auth.py b/trino/auth.py
@@ -89,6 +89,19 @@ def get_exceptions(self):
         except ImportError:
             raise RuntimeError("unable to import requests_kerberos")
 
+    def __eq__(self, other):
+        if not isinstance(other, KerberosAuthentication):
+            return False
+        return (self._config == other._config
+                and self._service_name == other._service_name
+                and self._mutual_authentication == other._mutual_authentication
+                and self._force_preemptive == other._force_preemptive
+                and self._hostname_override == other._hostname_override
+                and self._sanitize_mutual_error_response == other._sanitize_mutual_error_response
+                and self._principal == other._principal
+                and self._delegate == other._delegate
+                and self._ca_bundle == other._ca_bundle)
+
 
 class BasicAuthentication(Authentication):
     def __init__(self, username, password):
@@ -107,6 +120,11 @@ def set_http_session(self, http_session):
     def get_exceptions(self):
         return ()
 
+    def __eq__(self, other):
+        if not isinstance(other, BasicAuthentication):
+            return False
+        return self._username == other._username and self._password == other._password
+
 
 class _BearerAuth(AuthBase):
     """
@@ -132,6 +150,11 @@ def set_http_session(self, http_session):
     def get_exceptions(self):
         return ()
 
+    def __eq__(self, other):
+        if not isinstance(other, JWTAuthentication):
+            return False
+        return self.token == other.token
+
 
 def handle_redirect_auth_url(auth_url):
     print("Open the following URL in browser for the external authentication:")
@@ -239,3 +262,8 @@ def set_http_session(self, http_session):
 
     def get_exceptions(self):
         return ()
+
+    def __eq__(self, other):
+        if not isinstance(other, OAuth2Authentication):
+            return False
+        return self._redirect_auth_url == other._redirect_auth_url