fix: Update based on gemini-assit comments to make robust callback by handling async and removing encrypted_key complication

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

Author: agrawalradhika-cell

google/auth/aio/transport/mtls.py

@@ -17,9 +17,11 @@
 """
 
 import asyncio
+import inspect
 import logging
 from os import getenv, path
 
+from google.auth import exceptions
 import google.auth.transport._mtls_helper
 
 CERTIFICATE_CONFIGURATION_DEFAULT_PATH = "~/.config/gcloud/certificate_config.json"
@@ -71,8 +73,35 @@ def has_default_client_cert_source():
     return False
 
 
+async def default_client_cert_source():
+    """Get a callback which returns the default client SSL credentials.
+
+    Returns:
+        Callable[[], [bytes, bytes]]: A callback which returns the default
+            client certificate bytes and private key bytes, both in PEM format.
+
+    Raises:
+        google.auth.exceptions.DefaultClientCertSourceError: If the default
+            client SSL credentials don't exist or are malformed.
+    """
+    if not has_default_client_cert_source():
+        raise exceptions.MutualTLSChannelError(
+            "Default client cert source doesn't exist"
+        )
+
+    async def callback():
+        try:
+            _, cert_bytes, key_bytes = await get_client_cert_and_key()
+        except (OSError, RuntimeError, ValueError) as caught_exc:
+            new_exc = exceptions.MutualTLSChannelError(caught_exc)
+            raise new_exc from caught_exc
+
+        return cert_bytes, key_bytes
+
+    return callback
+
+
 async def get_client_ssl_credentials(
-    generate_encrypted_key=False,
     certificate_config_path=None,
 ):
     """Returns the client side certificate, private key and passphrase.
@@ -82,10 +111,6 @@ async def get_client_ssl_credentials(
                Currently, only X.509 workload certificates are supported.
 
     Args:
-        generate_encrypted_key (bool): If set to True, encrypted private key
-            and passphrase will be generated; otherwise, unencrypted private key
-            will be generated and passphrase will be None. This option only
-            affects keys obtained via context_aware_metadata.json.
         certificate_config_path (str): The certificate_config.json file path.
 
     Returns:
@@ -131,10 +156,12 @@ async def get_client_cert_and_key(client_cert_callback=None):
             the cert and key.
     """
     if client_cert_callback:
-        cert, key = client_cert_callback()
+        result = client_cert_callback()
+        if inspect.isawaitable(result):
+            cert, key = await result
+        else:
+            cert, key = result
         return True, cert, key
 
-    has_cert, cert, key, _ = await get_client_ssl_credentials(
-        generate_encrypted_key=False
-    )
+    has_cert, cert, key, _ = await get_client_ssl_credentials()
     return has_cert, cert, key

tests/transport/test_aio_mtls_helper.py

@@ -45,13 +45,71 @@ def test__check_config_path_not_found(self, mock_exists):
     @mock.patch("google.auth.aio.transport.mtls._check_config_path")
     @mock.patch("google.auth.aio.transport.mtls.getenv")
     def test_has_default_client_cert_source_env_var(self, mock_getenv, mock_check):
-        # Mocking so the default path fails but the env var path succeeds
         custom_path = "/custom/path.json"
         mock_check.side_effect = lambda x: custom_path if x == custom_path else None
         mock_getenv.return_value = custom_path
 
         assert mtls.has_default_client_cert_source() is True
 
+    @mock.patch("google.auth.aio.transport.mtls._check_config_path")
+    @mock.patch("google.auth.aio.transport.mtls.getenv")
+    def test_has_default_client_cert_source_check_priority(
+        self, mock_getenv, mock_check
+    ):
+        mock_check.return_value = "/default/path.json"
+
+        assert mtls.has_default_client_cert_source() is True
+        mock_getenv.assert_not_called()
+
+    @pytest.mark.asyncio
+    @mock.patch(
+        "google.auth.aio.transport.mtls.get_client_cert_and_key",
+        new_callable=mock.AsyncMock,
+    )
+    @mock.patch("google.auth.aio.transport.mtls.has_default_client_cert_source")
+    async def test_default_client_cert_source_success(
+        self, mock_has_default, mock_get_cert_key
+    ):
+        mock_has_default.return_value = True
+        mock_get_cert_key.return_value = (True, CERT_DATA, KEY_DATA)
+
+        callback = await mtls.default_client_cert_source()
+
+        cert, key = await callback()
+
+        assert cert == CERT_DATA
+        assert key == KEY_DATA
+        mock_has_default.assert_called_once()
+        mock_get_cert_key.assert_called_once()
+
+    @pytest.mark.asyncio
+    @mock.patch(
+        "google.auth.aio.transport.mtls.has_default_client_cert_source",
+        return_value=False,
+    )
+    async def test_default_client_cert_source_not_found(self, mock_has_default):
+        with pytest.raises(exceptions.MutualTLSChannelError, match="doesn't exist"):
+            await mtls.default_client_cert_source()
+
+    @pytest.mark.asyncio
+    @mock.patch(
+        "google.auth.aio.transport.mtls.get_client_cert_and_key",
+        new_callable=mock.AsyncMock,
+    )
+    @mock.patch(
+        "google.auth.aio.transport.mtls.has_default_client_cert_source",
+        return_value=True,
+    )
+    async def test_default_client_cert_source_callback_wraps_exception(
+        self, mock_has, mock_get
+    ):
+        mock_get.side_effect = ValueError("Format error")
+        callback = await mtls.default_client_cert_source()
+
+        with pytest.raises(exceptions.MutualTLSChannelError) as excinfo:
+            await callback()
+        assert "Format error" in str(excinfo.value)
+
     @pytest.mark.asyncio
     @mock.patch("google.auth.transport._mtls_helper._get_workload_cert_and_key")
     async def test_get_client_ssl_credentials_success(self, mock_workload):
@@ -64,6 +122,17 @@ async def test_get_client_ssl_credentials_success(self, mock_workload):
         assert key == KEY_DATA
         assert passphrase is None
 
+    @pytest.mark.asyncio
+    @mock.patch("google.auth.aio.transport.mtls.get_client_ssl_credentials")
+    async def test_get_client_cert_and_key_no_credentials_found(self, mock_get_ssl):
+        mock_get_ssl.return_value = (False, None, None, None)
+
+        success, cert, key = await mtls.get_client_cert_and_key(None)
+
+        assert success is False
+        assert cert is None
+        assert key is None
+
     @pytest.mark.asyncio
     async def test_get_client_cert_and_key_callback(self):
         # The callback should be tried first and return immediately
@@ -79,39 +148,33 @@ async def test_get_client_cert_and_key_callback(self):
     @pytest.mark.asyncio
     @mock.patch("google.auth.aio.transport.mtls.get_client_ssl_credentials")
     async def test_get_client_cert_and_key_default(self, mock_get_ssl):
-        # If no callback, it should call get_client_ssl_credentials
         mock_get_ssl.return_value = (True, CERT_DATA, KEY_DATA, None)
 
         success, cert, key = await mtls.get_client_cert_and_key(None)
 
         assert success is True
         assert cert == CERT_DATA
         assert key == KEY_DATA
-        mock_get_ssl.assert_called_with(generate_encrypted_key=False)
+        mock_get_ssl.assert_called_once()
 
     @pytest.mark.asyncio
     @mock.patch("google.auth.transport._mtls_helper._get_workload_cert_and_key")
     async def test_get_client_ssl_credentials_error(self, mock_workload):
-        """Tests that ClientCertError is propagated correctly."""
-        # Setup the mock to raise the specific google-auth exception
         mock_workload.side_effect = exceptions.ClientCertError(
             "Failed to read metadata"
         )
 
-        # Verify that calling our function raises the same exception
         with pytest.raises(exceptions.ClientCertError, match="Failed to read metadata"):
             await mtls.get_client_ssl_credentials()
 
     @pytest.mark.asyncio
     @mock.patch("google.auth.aio.transport.mtls.get_client_ssl_credentials")
     async def test_get_client_cert_and_key_exception_propagation(self, mock_get_ssl):
-        """Tests that get_client_cert_and_key propagates errors from its internal calls."""
         mock_get_ssl.side_effect = exceptions.ClientCertError(
             "Underlying credentials failed"
         )
 
         with pytest.raises(
             exceptions.ClientCertError, match="Underlying credentials failed"
         ):
-            # Pass None for callback so it attempts to call get_client_ssl_credentials
             await mtls.get_client_cert_and_key(client_cert_callback=None)