chore: Control each response for calls to the IAM Mock Server

oauth2_http/javatests/com/google/auth/oauth2/GoogleCredentialsTest.java

@@ -33,6 +33,7 @@
 
 import static org.junit.Assert.*;
 
+import com.google.api.client.http.HttpStatusCodes;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.util.Clock;
 import com.google.auth.Credentials;
@@ -609,6 +610,7 @@ public void fromStream_Impersonation_providesToken_WithQuotaProject() throws IOE
     transportFactory.transport.setExpireTime(ImpersonatedCredentialsTest.getDefaultExpireTime());
     transportFactory.transport.setAccessTokenEndpoint(
         ImpersonatedCredentialsTest.IMPERSONATION_URL);
+    transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
 
     InputStream impersonationCredentialsStream =
         ImpersonatedCredentialsTest.writeImpersonationCredentialsStream(
@@ -669,6 +671,7 @@ public void fromStream_Impersonation_providesToken_WithoutQuotaProject() throws
     transportFactory.transport.setExpireTime(ImpersonatedCredentialsTest.getDefaultExpireTime());
     transportFactory.transport.setAccessTokenEndpoint(
         ImpersonatedCredentialsTest.IMPERSONATION_URL);
+    transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
 
     InputStream impersonationCredentialsStream =
         ImpersonatedCredentialsTest.writeImpersonationCredentialsStream(

oauth2_http/javatests/com/google/auth/oauth2/IamUtilsTest.java

@@ -31,13 +31,15 @@
 package com.google.auth.oauth2;
 
 import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 import static org.junit.Assert.assertTrue;
 
 import com.google.api.client.http.HttpStatusCodes;
 import com.google.auth.ServiceAccountSigner;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
+import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -49,17 +51,53 @@ public class IamUtilsTest {
   private static final String CLIENT_EMAIL =
       "36680232662-vrd7ji19qe3nelgchd0ah2csanun6bnr@developer.gserviceaccount.com";
 
+  private ServiceAccountCredentials credentials;
+
+  @Before
+  public void setup() throws IOException {
+    // Mock this call for the Credentials because the IAM SignBlob RPC requires an access token. The
+    // call is initialized with HttpCredentialsAdapter which will make a call to get the access
+    // token
+    credentials = Mockito.mock(ServiceAccountCredentials.class);
+    Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
+  }
+
   @Test
-  public void sign_noRetry() throws IOException {
+  public void sign_success_noRetry() {
     byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
 
-    // Mock this call because signing requires an access token. The call is initialized with
-    // HttpCredentialsAdapter which will make a call to get the access token
-    ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
-    Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
+    ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
+        new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
+    transportFactory.transport.setSignedBlob(expectedSignature);
+    transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
+    transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
+
+    byte[] signature =
+        IamUtils.sign(
+            CLIENT_EMAIL,
+            credentials,
+            transportFactory.transport,
+            expectedSignature,
+            ImmutableMap.of());
+    assertArrayEquals(expectedSignature, signature);
+
+    assertEquals(1, transportFactory.transport.getNumRequests());
+  }
+
+  // The SignBlob RPC will retry up to three times before it gives up. This test will return two
+  // 5xx status codes before returning a success. This test covers the cases where the number of
+  // retry attempts is below the configured retry attempt count bounds (3 attempts).
+  @Test
+  public void sign_retryTwoTimes_success() {
+    byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
 
     ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
         new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE, "Unavailable");
+    transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
     transportFactory.transport.setSignedBlob(expectedSignature);
     transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
 
@@ -71,22 +109,97 @@ public void sign_noRetry() throws IOException {
             expectedSignature,
             ImmutableMap.of());
     assertArrayEquals(expectedSignature, signature);
+
+    // Expect that three requests are made (2 failures which are retries + 1 final requests which
+    // resulted in a successful response)
+    assertEquals(3, transportFactory.transport.getNumRequests());
   }
 
+  // The rpc will retry up to three times before it gives up. This test will enqueue three failed
+  // status codes + messages before returning a success. After the third retry attempt, the request
+  // will try one last time and the result will be reported back to the user.
   @Test
-  public void sign_4xxServerError_exception() throws IOException {
+  public void sign_retryThreeTimes_success() {
     byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
 
-    // Mock this call because signing requires an access token. The call is initialized with
-    // HttpCredentialsAdapter which will make a call to get the access token
-    ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
-    Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
+    ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
+        new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
+    transportFactory.transport.setSignedBlob(expectedSignature);
+    transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE, "Unavailable");
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_SERVER_ERROR, "Server Error");
+    transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
+
+    byte[] signature =
+        IamUtils.sign(
+            CLIENT_EMAIL,
+            credentials,
+            transportFactory.transport,
+            expectedSignature,
+            ImmutableMap.of());
+    assertArrayEquals(expectedSignature, signature);
+
+    // Expect that three requests are made (3 failures which are retried + 1 final request which
+    // resulted the final success response)
+    assertEquals(4, transportFactory.transport.getNumRequests());
+  }
+
+  // The rpc will retry up to three times before it gives up. This test will enqueue four failed
+  // status codes + messages before returning a success. After the third retry attempt, the request
+  // will try one last time and the result will be reported back to the user.
+  @Test
+  public void sign_retryThreeTimes_exception() {
+    byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
 
     ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
         new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
     transportFactory.transport.setSignedBlob(expectedSignature);
     transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
-    transportFactory.transport.setErrorResponseCodeAndMessage(
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE, "Unavailable");
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_SERVER_ERROR, "Server Error");
+    transportFactory.transport.addStatusCodeAndMessage(
+        HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
+    transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
+
+    ServiceAccountSigner.SigningException exception =
+        assertThrows(
+            ServiceAccountSigner.SigningException.class,
+            () ->
+                IamUtils.sign(
+                    CLIENT_EMAIL,
+                    credentials,
+                    transportFactory.transport,
+                    expectedSignature,
+                    ImmutableMap.of()));
+    assertTrue(exception.getMessage().contains("Failed to sign the provided bytes"));
+    assertTrue(
+        exception
+            .getCause()
+            .getMessage()
+            .contains("Unexpected Error code 502 trying to sign provided bytes"));
+
+    // Expect that three requests are made (3 failures which are retried + 1 final request which
+    // resulted in another failed response)
+    assertEquals(4, transportFactory.transport.getNumRequests());
+  }
+
+  @Test
+  public void sign_4xxError_noRetry_exception() {
+    byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
+
+    ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
+        new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
+    transportFactory.transport.setSignedBlob(expectedSignature);
+    transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
+    transportFactory.transport.addStatusCodeAndMessage(
         HttpStatusCodes.STATUS_CODE_UNAUTHORIZED, "Failed to sign the provided bytes");
 
     ServiceAccountSigner.SigningException exception =
@@ -105,5 +218,8 @@ public void sign_4xxServerError_exception() throws IOException {
             .getCause()
             .getMessage()
             .contains("Error code 401 trying to sign provided bytes:"));
+
+    // Only one request will have been made for a 4xx error (no retries)
+    assertEquals(1, transportFactory.transport.getNumRequests());
   }
 }