Intermediate changes

4 files changed, 170 insertions, 54 deletions

diff --git a/contrib/python/google-auth/py3/tests/compute_engine/test__metadata.py b/contrib/python/google-auth/py3/tests/compute_engine/test__metadata.py
index 352342f1509..a06dc4fa191 100644
--- a/contrib/python/google-auth/py3/tests/compute_engine/test__metadata.py
+++ b/contrib/python/google-auth/py3/tests/compute_engine/test__metadata.py
@@ -433,6 +433,74 @@ def test_get_universe_domain_not_found():
     assert universe_domain == "googleapis.com"
 
 
+def test_get_universe_domain_retryable_error_failure():
+    # Test that if the universe domain endpoint returns a retryable error
+    # we should retry.
+    #
+    # In this case, the error persists, and we still fail after retrying.
+    request = make_request("too many requests", status=http_client.TOO_MANY_REQUESTS)
+
+    with pytest.raises(exceptions.TransportError) as excinfo:
+        _metadata.get_universe_domain(request)
+
+    assert excinfo.match(r"Compute Engine Metadata server unavailable")
+
+    request.assert_called_with(
+        method="GET",
+        url=_metadata._METADATA_ROOT + "universe/universe_domain",
+        headers=_metadata._METADATA_HEADERS,
+    )
+    assert request.call_count == 5
+
+
+def test_get_universe_domain_retryable_error_success():
+    # Test that if the universe domain endpoint returns a retryable error
+    # we should retry.
+    #
+    # In this case, the error is temporary, and we succeed after retrying.
+    request_error = make_request(
+        "too many requests", status=http_client.TOO_MANY_REQUESTS
+    )
+    request_ok = make_request(
+        "fake_universe_domain", headers={"content-type": "text/plain"}
+    )
+
+    class _RequestErrorOnce:
+        """This class forwards the request parameters to `request_error` once.
+
+        All subsequent calls are forwarded to `request_ok`.
+        """
+
+        def __init__(self, request_error, request_ok):
+            self._request_error = request_error
+            self._request_ok = request_ok
+            self._call_index = 0
+
+        def request(self, *args, **kwargs):
+            if self._call_index == 0:
+                self._call_index += 1
+                return self._request_error(*args, **kwargs)
+
+            return self._request_ok(*args, **kwargs)
+
+    request = _RequestErrorOnce(request_error, request_ok).request
+
+    universe_domain = _metadata.get_universe_domain(request)
+
+    request_error.assert_called_once_with(
+        method="GET",
+        url=_metadata._METADATA_ROOT + "universe/universe_domain",
+        headers=_metadata._METADATA_HEADERS,
+    )
+    request_ok.assert_called_once_with(
+        method="GET",
+        url=_metadata._METADATA_ROOT + "universe/universe_domain",
+        headers=_metadata._METADATA_HEADERS,
+    )
+
+    assert universe_domain == "fake_universe_domain"
+
+
 def test_get_universe_domain_other_error():
     # Test that if the universe domain endpoint returns an error other than 404
     # we should throw the error
diff --git a/contrib/python/google-auth/py3/tests/transport/test__mtls_helper.py b/contrib/python/google-auth/py3/tests/transport/test__mtls_helper.py
index b195616dd57..f6e20b726a8 100644
--- a/contrib/python/google-auth/py3/tests/transport/test__mtls_helper.py
+++ b/contrib/python/google-auth/py3/tests/transport/test__mtls_helper.py
@@ -111,15 +111,15 @@ class TestCertAndKeyRegex(object):
         )
 
 
-class TestCheckaMetadataPath(object):
+class TestCheckConfigPath(object):
     def test_success(self):
         metadata_path = os.path.join(pytest.data_dir, "context_aware_metadata.json")
-        returned_path = _mtls_helper._check_dca_metadata_path(metadata_path)
+        returned_path = _mtls_helper._check_config_path(metadata_path)
         assert returned_path is not None
 
     def test_failure(self):
         metadata_path = os.path.join(pytest.data_dir, "not_exists.json")
-        returned_path = _mtls_helper._check_dca_metadata_path(metadata_path)
+        returned_path = _mtls_helper._check_config_path(metadata_path)
         assert returned_path is None
 
 
@@ -275,21 +275,24 @@ class TestRunCertProviderCommand(object):
 
 class TestGetClientSslCredentials(object):
     @mock.patch(
-        "google.auth.transport._mtls_helper._run_cert_provider_command", autospec=True
+        "google.auth.transport._mtls_helper._get_workload_cert_and_key", autospec=True
     )
-    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
     @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
+        "google.auth.transport._mtls_helper._run_cert_provider_command", autospec=True
     )
-    def test_success(
+    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
+    def test_success_with_context_aware_metadata(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_run_cert_provider_command,
+        mock_get_workload_cert_and_key,
     ):
-        mock_check_dca_metadata_path.return_value = True
+        mock_check_config_path.return_value = "/path/to/config"
         mock_load_json_file.return_value = {"cert_provider_command": ["command"]}
         mock_run_cert_provider_command.return_value = (b"cert", b"key", None)
+        mock_get_workload_cert_and_key.return_value = (None, None)
         has_cert, cert, key, passphrase = _mtls_helper.get_client_ssl_credentials()
         assert has_cert
         assert cert == b"cert"
@@ -297,10 +300,42 @@ class TestGetClientSslCredentials(object):
         assert passphrase is None
 
     @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
+        "google.auth.transport._mtls_helper._read_cert_and_key_files", autospec=True
     )
-    def test_success_without_metadata(self, mock_check_dca_metadata_path):
-        mock_check_dca_metadata_path.return_value = False
+    @mock.patch(
+        "google.auth.transport._mtls_helper._get_cert_config_path", autospec=True
+    )
+    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
+    def test_success_with_certificate_config(
+        self,
+        mock_check_config_path,
+        mock_load_json_file,
+        mock_get_cert_config_path,
+        mock_read_cert_and_key_files,
+    ):
+        cert_config_path = "/path/to/config"
+        mock_check_config_path.return_value = cert_config_path
+        mock_load_json_file.return_value = {
+            "cert_configs": {
+                "workload": {"cert_path": "cert/path", "key_path": "key/path"}
+            }
+        }
+        mock_get_cert_config_path.return_value = cert_config_path
+        mock_read_cert_and_key_files.return_value = (
+            pytest.public_cert_bytes,
+            pytest.private_key_bytes,
+        )
+
+        has_cert, cert, key, passphrase = _mtls_helper.get_client_ssl_credentials()
+        assert has_cert
+        assert cert == pytest.public_cert_bytes
+        assert key == pytest.private_key_bytes
+        assert passphrase is None
+
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
+    def test_success_without_metadata(self, mock_check_config_path):
+        mock_check_config_path.return_value = False
         has_cert, cert, key, passphrase = _mtls_helper.get_client_ssl_credentials()
         assert not has_cert
         assert cert is None
@@ -308,21 +343,24 @@ class TestGetClientSslCredentials(object):
         assert passphrase is None
 
     @mock.patch(
-        "google.auth.transport._mtls_helper._run_cert_provider_command", autospec=True
+        "google.auth.transport._mtls_helper._get_workload_cert_and_key", autospec=True
     )
-    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
     @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
+        "google.auth.transport._mtls_helper._run_cert_provider_command", autospec=True
     )
+    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
     def test_success_with_encrypted_key(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_run_cert_provider_command,
+        mock_get_workload_cert_and_key,
     ):
-        mock_check_dca_metadata_path.return_value = True
+        mock_check_config_path.return_value = "/path/to/config"
         mock_load_json_file.return_value = {"cert_provider_command": ["command"]}
         mock_run_cert_provider_command.return_value = (b"cert", b"key", b"passphrase")
+        mock_get_workload_cert_and_key.return_value = (None, None)
         has_cert, cert, key, passphrase = _mtls_helper.get_client_ssl_credentials(
             generate_encrypted_key=True
         )
@@ -334,15 +372,20 @@ class TestGetClientSslCredentials(object):
             ["command", "--with_passphrase"], expect_encrypted_key=True
         )
 
-    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
     @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
+        "google.auth.transport._mtls_helper._get_workload_cert_and_key", autospec=True
     )
+    @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
     def test_missing_cert_command(
-        self, mock_check_dca_metadata_path, mock_load_json_file
+        self,
+        mock_check_config_path,
+        mock_load_json_file,
+        mock_get_workload_cert_and_key,
     ):
-        mock_check_dca_metadata_path.return_value = True
+        mock_check_config_path.return_value = "/path/to/config"
         mock_load_json_file.return_value = {}
+        mock_get_workload_cert_and_key.return_value = (None, None)
         with pytest.raises(exceptions.ClientCertError):
             _mtls_helper.get_client_ssl_credentials()
 
@@ -350,17 +393,15 @@ class TestGetClientSslCredentials(object):
         "google.auth.transport._mtls_helper._run_cert_provider_command", autospec=True
     )
     @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
-    @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
-    )
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
     def test_customize_context_aware_metadata_path(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_run_cert_provider_command,
     ):
         context_aware_metadata_path = "/path/to/metata/data"
-        mock_check_dca_metadata_path.return_value = context_aware_metadata_path
+        mock_check_config_path.return_value = context_aware_metadata_path
         mock_load_json_file.return_value = {"cert_provider_command": ["command"]}
         mock_run_cert_provider_command.return_value = (b"cert", b"key", None)
 
@@ -372,7 +413,7 @@ class TestGetClientSslCredentials(object):
         assert cert == b"cert"
         assert key == b"key"
         assert passphrase is None
-        mock_check_dca_metadata_path.assert_called_with(context_aware_metadata_path)
+        mock_check_config_path.assert_called_with(context_aware_metadata_path)
         mock_load_json_file.assert_called_with(context_aware_metadata_path)
 
 
@@ -520,7 +561,7 @@ class TestGetCertConfigPath(object):
         mock_path_exists.return_value = True
         returned_path = _mtls_helper._get_cert_config_path()
         expected_path = os.path.expanduser(
-            _mtls_helper._CERTIFICATE_CONFIGURATION_DEFAULT_PATH
+            _mtls_helper.CERTIFICATE_CONFIGURATION_DEFAULT_PATH
         )
         assert returned_path == expected_path
 
diff --git a/contrib/python/google-auth/py3/tests/transport/test_grpc.py b/contrib/python/google-auth/py3/tests/transport/test_grpc.py
index 9badb59b284..80b24e86d41 100644
--- a/contrib/python/google-auth/py3/tests/transport/test_grpc.py
+++ b/contrib/python/google-auth/py3/tests/transport/test_grpc.py
@@ -143,12 +143,10 @@ class TestAuthMetadataPlugin(object):
 @mock.patch("grpc.secure_channel", autospec=True)
 class TestSecureAuthorizedChannel(object):
     @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
-    @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
-    )
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
     def test_secure_authorized_channel_adc(
         self,
-        check_dca_metadata_path,
+        check_config_path,
         load_json_file,
         secure_channel,
         ssl_channel_credentials,
@@ -162,7 +160,7 @@ class TestSecureAuthorizedChannel(object):
 
         # Mock the context aware metadata and client cert/key so mTLS SSL channel
         # will be used.
-        check_dca_metadata_path.return_value = METADATA_PATH
+        check_config_path.return_value = METADATA_PATH
         load_json_file.return_value = {"cert_provider_command": ["some command"]}
         get_client_ssl_credentials.return_value = (
             True,
@@ -332,12 +330,10 @@ class TestSecureAuthorizedChannel(object):
         )
 
     @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
-    @mock.patch(
-        "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
-    )
+    @mock.patch("google.auth.transport._mtls_helper._check_config_path", autospec=True)
     def test_secure_authorized_channel_with_client_cert_callback_failure(
         self,
-        check_dca_metadata_path,
+        check_config_path,
         load_json_file,
         secure_channel,
         ssl_channel_credentials,
@@ -401,19 +397,17 @@ class TestSecureAuthorizedChannel(object):
     "google.auth.transport._mtls_helper.get_client_ssl_credentials", autospec=True
 )
 @mock.patch("google.auth.transport._mtls_helper._load_json_file", autospec=True)
[email protected](
-    "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
-)
[email protected]("google.auth.transport._mtls_helper._check_config_path", autospec=True)
 class TestSslCredentials(object):
     def test_no_context_aware_metadata(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_get_client_ssl_credentials,
         mock_ssl_channel_credentials,
     ):
         # Mock that the metadata file doesn't exist.
-        mock_check_dca_metadata_path.return_value = None
+        mock_check_config_path.return_value = None
 
         with mock.patch.dict(
             os.environ, {environment_vars.GOOGLE_API_USE_CLIENT_CERTIFICATE: "true"}
@@ -430,12 +424,12 @@ class TestSslCredentials(object):
 
     def test_get_client_ssl_credentials_failure(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_get_client_ssl_credentials,
         mock_ssl_channel_credentials,
     ):
-        mock_check_dca_metadata_path.return_value = METADATA_PATH
+        mock_check_config_path.return_value = METADATA_PATH
         mock_load_json_file.return_value = {"cert_provider_command": ["some command"]}
 
         # Mock that client cert and key are not loaded and exception is raised.
@@ -449,12 +443,12 @@ class TestSslCredentials(object):
 
     def test_get_client_ssl_credentials_success(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_get_client_ssl_credentials,
         mock_ssl_channel_credentials,
     ):
-        mock_check_dca_metadata_path.return_value = METADATA_PATH
+        mock_check_config_path.return_value = METADATA_PATH
         mock_load_json_file.return_value = {"cert_provider_command": ["some command"]}
         mock_get_client_ssl_credentials.return_value = (
             True,
@@ -477,7 +471,7 @@ class TestSslCredentials(object):
 
     def test_get_client_ssl_credentials_without_client_cert_env(
         self,
-        mock_check_dca_metadata_path,
+        mock_check_config_path,
         mock_load_json_file,
         mock_get_client_ssl_credentials,
         mock_ssl_channel_credentials,
@@ -487,7 +481,7 @@ class TestSslCredentials(object):
 
         assert ssl_credentials.ssl_credentials is not None
         assert not ssl_credentials.is_mtls
-        mock_check_dca_metadata_path.assert_not_called()
+        mock_check_config_path.assert_not_called()
         mock_load_json_file.assert_not_called()
         mock_get_client_ssl_credentials.assert_not_called()
         mock_ssl_channel_credentials.assert_called_once()
diff --git a/contrib/python/google-auth/py3/tests/transport/test_mtls.py b/contrib/python/google-auth/py3/tests/transport/test_mtls.py
index b62063e4797..ea549ae142b 100644
--- a/contrib/python/google-auth/py3/tests/transport/test_mtls.py
+++ b/contrib/python/google-auth/py3/tests/transport/test_mtls.py
@@ -16,17 +16,30 @@ import mock
 import pytest  # type: ignore
 
 from google.auth import exceptions
+from google.auth.transport import _mtls_helper
 from google.auth.transport import mtls
 
 
[email protected](
-    "google.auth.transport._mtls_helper._check_dca_metadata_path", autospec=True
-)
-def test_has_default_client_cert_source(check_dca_metadata_path):
-    check_dca_metadata_path.return_value = mock.Mock()
[email protected]("google.auth.transport._mtls_helper._check_config_path", autospec=True)
+def test_has_default_client_cert_source(check_config_path):
+    def return_path_for_metadata(path):
+        return mock.Mock() if path == _mtls_helper.CONTEXT_AWARE_METADATA_PATH else None
+
+    check_config_path.side_effect = return_path_for_metadata
+    assert mtls.has_default_client_cert_source()
+
+    def return_path_for_cert_config(path):
+        return (
+            mock.Mock()
+            if path == _mtls_helper.CERTIFICATE_CONFIGURATION_DEFAULT_PATH
+            else None
+        )
+
+    check_config_path.side_effect = return_path_for_cert_config
     assert mtls.has_default_client_cert_source()
 
-    check_dca_metadata_path.return_value = None
+    check_config_path.side_effect = None
+    check_config_path.return_value = None
     assert not mtls.has_default_client_cert_source()