feat(auth): Add unit tests for auth failures in gax layer

src/auth/src/errors.rs

@@ -63,7 +63,7 @@ impl CredentialError {
     /// # Arguments
     /// * `is_retryable` - A boolean indicating whether the error is retryable.
     /// * `message` - The underlying error that caused the auth failure.
-    pub(crate) fn from_str<T: Into<String>>(is_retryable: bool, message: T) -> Self {
+    pub fn from_str<T: Into<String>>(is_retryable: bool, message: T) -> Self {
         CredentialError::new(
             is_retryable,
             CredentialErrorImpl::SimpleMessage(message.into()),

src/gax/src/http_client/mod.rs

@@ -76,14 +76,6 @@ impl ReqwestClient {
         body: Option<I>,
         options: crate::options::RequestOptions,
     ) -> Result<O> {
-        let auth_headers = self
-            .cred
-            .get_headers()
-            .await
-            .map_err(Error::authentication)?;
-        for header in auth_headers.into_iter() {
-            builder = builder.header(header.0, header.1);
-        }
         if let Some(user_agent) = options.user_agent() {
             builder = builder.header(
                 reqwest::header::USER_AGENT,
@@ -186,6 +178,14 @@ impl ReqwestClient {
         {
             builder = builder.timeout(timeout);
         }
+        let auth_headers = self
+            .cred
+            .get_headers()
+            .await
+            .map_err(Error::authentication)?;
+        for header in auth_headers.into_iter() {
+            builder = builder.header(header.0, header.1);
+        }
         let response = builder.send().await.map_err(Error::io)?;
         if !response.status().is_success() {
             return Self::to_http_error(response).await;

src/gax/src/retry_policy.rs

@@ -38,6 +38,8 @@
 //!
 //! [idempotent]: https://en.wikipedia.org/wiki/Idempotence
 
+use auth::errors::CredentialError;
+
 use crate::error::Error;
 use crate::loop_state::LoopState;
 use std::sync::Arc;
@@ -234,9 +236,15 @@ impl RetryPolicy for Aip194Strict {
                 }
             }
             ErrorKind::Authentication => {
-                // This indicates the operation never left the client, so it
-                // safe to retry
-                LoopState::Continue(error)
+                if let Some(cred_err) = error.as_inner::<CredentialError>() {
+                    if cred_err.is_retryable() {
+                        LoopState::Continue(error)
+                    } else {
+                        LoopState::Permanent(error)
+                    }
+                } else {
+                    LoopState::Continue(error)
+                }
             }
             ErrorKind::Serde => LoopState::Permanent(error),
             ErrorKind::Other => LoopState::Permanent(error),

src/gax/tests/auth.rs

@@ -18,6 +18,7 @@ use auth::token::Token;
 use gax::http_client::ReqwestClient;
 use gax::options::*;
 use google_cloud_gax as gax;
+use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicyExt};
 use http::header::{HeaderName, HeaderValue};
 use serde_json::json;
 
@@ -75,6 +76,77 @@ async fn test_auth_headers() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn auth_error_retryable() -> Result<()> {
+    let (endpoint, _server) = echo_server::start().await?;
+    let retry_count = 3;
+    let mut mock = MockCredential::new();
+    mock.expect_get_headers()
+        .times(retry_count..)
+        .returning(|| Err(CredentialError::from_str(true, "mock retryable error")));
+
+    let retry_policy = Aip194Strict.with_attempt_limit(retry_count as u32);
+    let config = ClientConfig::default()
+        .set_credential(Credential::from(mock))
+        .set_retry_policy(retry_policy);
+    let client = ReqwestClient::new(config, &endpoint).await?;
+
+    let builder = client.builder(reqwest::Method::GET, "/echo".into());
+    let body = json!({});
+    let options = RequestOptions::default();
+    let result: std::result::Result<serde_json::Value, google_cloud_gax::error::Error> =
+        client.execute(builder, Some(body), options).await;
+
+    assert!(result.is_err());
+
+    if let Err(e) = result {
+        if let Some(cred_err) = e.as_inner::<CredentialError>() {
+            assert!(
+                cred_err.is_retryable(),
+                "Expected a retryable CredentialError, but got non-retryable"
+            );
+        } else {
+            panic!("Expected a CredentialError, but got some other error: {e:?}");
+        }
+    }
+
+    Ok(())
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn auth_error_non_retryable() -> Result<()> {
+    let (endpoint, _server) = echo_server::start().await?;
+    let mut mock = MockCredential::new();
+    mock.expect_get_headers()
+        .times(1)
+        .returning(|| Err(CredentialError::from_str(false, "mock non-retryable error")));
+
+    let config = ClientConfig::default().set_credential(Credential::from(mock));
+    let client = ReqwestClient::new(config, &endpoint).await?;
+
+    let builder = client.builder(reqwest::Method::GET, "/echo".into());
+    let body = serde_json::json!({});
+
+    let result: std::result::Result<serde_json::Value, google_cloud_gax::error::Error> = client
+        .execute(builder, Some(body), RequestOptions::default())
+        .await;
+
+    assert!(result.is_err());
+
+    if let Err(e) = result {
+        if let Some(cred_err) = e.as_inner::<CredentialError>() {
+            assert!(
+                !cred_err.is_retryable(),
+                "Expected a non-retryable CredentialError, but got retryable"
+            );
+        } else {
+            panic!("Expected a CredentialError, but got another error type: {e:?}");
+        }
+    }
+
+    Ok(())
+}
+
 fn get_header_value(response: &serde_json::Value, name: &str) -> Option<String> {
     response
         .as_object()