fix: Fix lint and unit tests

Signed-off-by: Radhika Agrawal <agrawalradhika@google.com>

Author: agrawalradhika-cell

packages/google-auth/google/auth/transport/grpc.py

@@ -24,7 +24,7 @@
 
 
 from google.auth import exceptions
-from google.auth import transport 
+from google.auth import transport
 from google.auth.transport import _mtls_helper
 from google.auth.transport import mtls
 from google.oauth2 import service_account
@@ -134,6 +134,8 @@ def secure_authorized_channel(
     target,
     ssl_credentials=None,
     client_cert_callback=None,
+    *,
+    auto_refresh: bool = True,
     **kwargs
 ):
     """Creates a secure authorized gRPC channel.
@@ -288,7 +290,8 @@ def my_client_cert_callback():
         )
 
     # If SSL credentials are not explicitly set, try client_cert_callback and ADC.
-    cached_cert = None
+    import typing
+    cached_cert: typing.Optional[bytes] = None
     if not ssl_credentials:
         use_client_cert = _mtls_helper.check_use_client_cert()
         if use_client_cert and client_cert_callback:
@@ -310,25 +313,21 @@ def my_client_cert_callback():
     composite_credentials = grpc.composite_channel_credentials(
         ssl_credentials, google_auth_credentials
     )
-    is_recreation = kwargs.pop("_is_recreation", False)
     channel = grpc.secure_channel(target, composite_credentials, **kwargs)
-    # Avoid wrapping if mTLS is disabled or if this is a channel recreation call
-    if cached_cert and not is_recreation:
-        # Package arguments so the channel can be recreated later
-        factory_args = {
-            "credentials": credentials,
-            "request": request,
-            "target": target,
-            "ssl_credentials": None,
-            "client_cert_callback": client_cert_callback,
-            "_is_recreation": True,  # Hidden flag to stop recursion
+    # Avoid wrapping if mTLS is disabled or if auto_refresh is False
+    if cached_cert and auto_refresh:
+        import functools
+        channel_fn = functools.partial(
+            secure_authorized_channel,
+            credentials,
+            request,
+            target,
+            client_cert_callback=None,
+            auto_refresh=False,
             **kwargs,
-        }
-        interceptor = _MTLSCallInterceptor()
-
-        wrapper = _MTLSRefreshingChannel(target, factory_args, channel, cached_cert)
-
-        interceptor._wrapper = wrapper
+        )
+        wrapper = MTLSRefreshingChannel(target, channel_fn, channel, cached_cert)
+        interceptor = CertRotationInterceptor(wrapper=wrapper)
         return grpc.intercept_channel(wrapper, interceptor)
     return channel
 
@@ -354,7 +353,8 @@ class SslCredentials:
 
     def __init__(self):
         use_client_cert = _mtls_helper.check_use_client_cert()
-        self._cached_cert = None
+        self._import typing
+    cached_cert: typing.Optional[bytes] = None
         if not use_client_cert:
             self._is_mtls = False
         else:
@@ -401,37 +401,39 @@ def is_mtls(self):
         return self._is_mtls
 
 
-class _MTLSCallInterceptor(
+class CertRotationInterceptor(
     grpc.UnaryUnaryClientInterceptor,
     grpc.UnaryStreamClientInterceptor,
     grpc.StreamUnaryClientInterceptor,
     grpc.StreamStreamClientInterceptor,
 ):
     """A gRPC client interceptor that provides automatic retry logic for mTLS certificate rotation.
-    
-    This interceptor wraps all gRPC client calls (unary and streaming) with retryable 
-    futures or iterators. Its primary role is to monitor responses for `UNAUTHENTICATED` 
-    errors. When an authentication failure occurs, it uses `_should_retry()` to check 
-    if a new mTLS certificate is available. If a new certificate is found, it signals 
-    its associated `_MTLSRefreshingChannel` wrapper to refresh the underlying gRPC 
+
+    This interceptor wraps all gRPC client calls (unary and streaming) with retryable
+    futures or iterators. Its primary role is to monitor responses for `UNAUTHENTICATED`
+    errors. When an authentication failure occurs, it uses `_should_retry()` to check
+    if a new mTLS certificate is available. If a new certificate is found, it signals
+    its associated `MTLSRefreshingChannel` wrapper to refresh the underlying gRPC
     channel's credentials and automatically replays the failed RPC.
     """
 
-    def __init__(self):
-        self._wrapper = None
+    def __init__(self, wrapper=None):
+        self._wrapper = wrapper
         self._max_retries = transport.DEFAULT_MAX_REFRESH_ATTEMPTS
 
     def _should_retry(self, code, retry_count, attempt_cert):
         """Determines if the RPC should be retried due to a certificate rotation.
-        
-         Returns a tuple: (should_retry, call_cert_bytes, call_key_bytes, passphrase).
-         """
+
+        Returns a tuple: (should_retry, call_cert_bytes, call_key_bytes, passphrase).
+        """
         if code != grpc.StatusCode.UNAUTHENTICATED or not self._wrapper:
             return False, None, None, None
 
         if retry_count >= self._max_retries:
             _LOGGER.debug(
-                "Max retries reached (%d/%d) for channel recreation.", retry_count, self._max_retries
+                "Max retries reached (%d/%d) for channel recreation.",
+                retry_count,
+                self._max_retries,
             )
             return False, None, None, None
 
@@ -483,39 +485,37 @@ def intercept_stream_stream(
         )
 
 
-class _MTLSRefreshingChannel(grpc.Channel):
-    def __init__(self, target, factory_args, initial_channel, initial_cert):
+class MTLSRefreshingChannel(grpc.Channel):
+    def __init__(self, target, channel_fn, initial_channel, initial_cert):
         self._target = target
-        self._factory_args = factory_args
+        self._channel_fn = channel_fn
         self._channel = initial_channel
         self._cached_cert = initial_cert
         self._lock = threading.Lock()
         self._subscribers = set()
 
-    def refresh_logic(self, count, call_cert_bytes=None, call_key_bytes=None, passphrase=None):
+    def refresh_logic(
+        self, count, call_cert_bytes=None, call_key_bytes=None, passphrase=None
+    ):
         with self._lock:
             if not call_cert_bytes or self._cached_cert == call_cert_bytes:
                 return
 
-            _LOGGER.debug(
-                "Wrapper: Refreshing mTLS channel. Retry count: %d", count
-            )
+            _LOGGER.debug("Wrapper: Refreshing mTLS channel. Retry count: %d", count)
             old_channel = self._channel
 
             if passphrase is not None:
                 call_key_bytes = _mtls_helper.decrypt_private_key(
                     call_key_bytes, passphrase
                 )
 
-            # The factory args must use the new credentials exactly to build the rotation channel
-            factory_args = self._factory_args.copy()
-            factory_args["client_cert_callback"] = None
-            factory_args["ssl_credentials"] = grpc.ssl_channel_credentials(
+            new_ssl_credentials = grpc.ssl_channel_credentials(
                 certificate_chain=call_cert_bytes,
                 private_key=call_key_bytes,
             )
-
-            self._channel = secure_authorized_channel(**factory_args)
+            self._channel = self._channel_fn(
+                ssl_credentials=new_ssl_credentials
+            )
             self._cached_cert = call_cert_bytes
             for callback in self._subscribers:
                 try:
@@ -630,15 +630,19 @@ def __next__(self):
     ("method", "timeout", "metadata", "credentials", "wait_for_ready"),
 )
 
+
 class _DeadlineExceededError(grpc.RpcError, grpc.Call):
     def __init__(self, details):
         super().__init__()
         self._details = details
+
     def code(self):
         return grpc.StatusCode.DEADLINE_EXCEEDED
+
     def details(self):
         return self._details
 
+
 class _RetryableUnaryResponseFuture(grpc.Future, grpc.Call):
     def __init__(
         self,
@@ -698,7 +702,9 @@ def _start_call(self):
                 elapsed = time.monotonic() - self._start_time
                 remaining = self._initial_timeout - elapsed
                 if remaining <= 0:
-                    raise _DeadlineExceededError("Deadline Exceeded during retry resolution.")
+                    raise _DeadlineExceededError(
+                        "Deadline Exceeded during retry resolution."
+                    )
                 call_details = _ClientCallDetails(
                     method=call_details.method,
                     timeout=remaining,
@@ -742,7 +748,9 @@ def _on_inner_future_done(self, inner_future):
             if can_replay and should_retry:
                 if getattr(self._interceptor, "_wrapper", None):
                     try:
-                        self._interceptor._wrapper.refresh_logic(1, call_cert, call_key, pwd)
+                        self._interceptor._wrapper.refresh_logic(
+                            1, call_cert, call_key, pwd
+                        )
                     except Exception as e:
                         with self._lock:
                             self._terminal_exception = e
@@ -757,18 +765,23 @@ def _on_inner_future_done(self, inner_future):
                     self._terminal_exception = e
 
             if self._interceptor._wrapper:
-                chk_should_retry, chk_cert, chk_key, chk_pwd = self._interceptor._should_retry(
-                    status_code, 0, self._attempt_cert
-                )
+                (
+                    chk_should_retry,
+                    chk_cert,
+                    chk_key,
+                    chk_pwd,
+                ) = self._interceptor._should_retry(status_code, 0, self._attempt_cert)
                 if chk_should_retry:
                     try:
-                        self._interceptor._wrapper.refresh_logic(1, chk_cert, chk_key, chk_pwd)
+                        self._interceptor._wrapper.refresh_logic(
+                            1, chk_cert, chk_key, chk_pwd
+                        )
                     except Exception:
                         pass
         with self._lock:
             self._completion_event.set()
             callbacks_to_fire = list(self._done_callbacks)
-            
+
         for fn in callbacks_to_fire:
             try:
                 fn(self)
@@ -929,7 +942,9 @@ def _start_call(self):
                 elapsed = time.monotonic() - self._start_time
                 remaining = self._initial_timeout - elapsed
                 if remaining <= 0:
-                    raise _DeadlineExceededError("Deadline Exceeded during retry resolution.")
+                    raise _DeadlineExceededError(
+                        "Deadline Exceeded during retry resolution."
+                    )
                 call_details = _ClientCallDetails(
                     method=call_details.method,
                     timeout=remaining,
@@ -959,9 +974,12 @@ def _on_inner_call_done(self, inner_call):
             if self._call is not inner_call:
                 return
             # Intercept and suppress premature callbacks for UNAUTHENTICATED.
-            # __next__ inherently handles this error and manages triggering callbacks 
+            # __next__ inherently handles this error and manages triggering callbacks
             # later if retriies are exhausted.
-            if callable(getattr(inner_call, "code", None)) and inner_call.code() == grpc.StatusCode.UNAUTHENTICATED:
+            if (
+                callable(getattr(inner_call, "code", None))
+                and inner_call.code() == grpc.StatusCode.UNAUTHENTICATED
+            ):
                 return
         self._trigger_callbacks()
 
@@ -994,42 +1012,52 @@ def __next__(self):
                     )
                 )
 
-                should_retry, call_cert, call_key, pwd = self._interceptor._should_retry(
+                (
+                    should_retry,
+                    call_cert,
+                    call_key,
+                    pwd,
+                ) = self._interceptor._should_retry(
                     status_code,
                     self._retry_count,
                     getattr(self, "_attempt_cert", None),
                 )
 
-                if (
-                    not self._yielded_any_response
-                    and can_replay
-                    and should_retry
-                ):
+                if not self._yielded_any_response and can_replay and should_retry:
                     try:
                         if getattr(self._interceptor, "_wrapper", None):
-                            self._interceptor._wrapper.refresh_logic(1, call_cert, call_key, pwd)
-    
+                            self._interceptor._wrapper.refresh_logic(
+                                1, call_cert, call_key, pwd
+                            )
+
                         with self._lock:
                             self._retry_count += 1
                             self._start_call()
-                            
+
                     except Exception as fallback_e:
                         self._trigger_callbacks()
                         raise fallback_e
-                        
+
                     continue
                 else:
                     # Non-retryable error, check if another rotation happened while we were finishing
                     if getattr(self._interceptor, "_wrapper", None):
-                        chk_should_retry, chk_cert, chk_key, chk_pwd = self._interceptor._should_retry(
+                        (
+                            chk_should_retry,
+                            chk_cert,
+                            chk_key,
+                            chk_pwd,
+                        ) = self._interceptor._should_retry(
                             status_code, 0, getattr(self, "_attempt_cert", None)
                         )
                         if chk_should_retry:
                             try:
-                                self._interceptor._wrapper.refresh_logic(1, chk_cert, chk_key, chk_pwd)
+                                self._interceptor._wrapper.refresh_logic(
+                                    1, chk_cert, chk_key, chk_pwd
+                                )
                             except Exception:
-                                pass # Terminal anyway
-                                
+                                pass  # Terminal anyway
+
                     self._trigger_callbacks()
                     raise e