Skip to content

Commit c79f707

Browse files
lqiu96lqiu96
authored
chore: Control each response for calls to the IAM Mock Server (#1455)
* fix: Retry sign blob call with exponential backoff * chore: Add header for IamUtilsTest * chore: Add comments for IamUtilsTest * chore: Control each response for calls to the IAM Mock Server * chore: Move IAM Retry status codes to IamUtils * chore: Address PR comments * chore: Address PR comments * chore: Address PR comments
1 parent d2a5833 commit c79f707

4 files changed

Lines changed: 210 additions & 84 deletions

File tree

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333

3434
import static org.junit.Assert.*;
3535

36+
import com.google.api.client.http.HttpStatusCodes;
3637
import com.google.api.client.json.GenericJson;
3738
import com.google.api.client.util.Clock;
3839
import com.google.auth.Credentials;
@@ -609,6 +610,7 @@ public void fromStream_Impersonation_providesToken_WithQuotaProject() throws IOE
609610
transportFactory.transport.setExpireTime(ImpersonatedCredentialsTest.getDefaultExpireTime());
610611
transportFactory.transport.setAccessTokenEndpoint(
611612
ImpersonatedCredentialsTest.IMPERSONATION_URL);
613+
transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
612614

613615
InputStream impersonationCredentialsStream =
614616
ImpersonatedCredentialsTest.writeImpersonationCredentialsStream(
@@ -669,6 +671,7 @@ public void fromStream_Impersonation_providesToken_WithoutQuotaProject() throws
669671
transportFactory.transport.setExpireTime(ImpersonatedCredentialsTest.getDefaultExpireTime());
670672
transportFactory.transport.setAccessTokenEndpoint(
671673
ImpersonatedCredentialsTest.IMPERSONATION_URL);
674+
transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
672675

673676
InputStream impersonationCredentialsStream =
674677
ImpersonatedCredentialsTest.writeImpersonationCredentialsStream(

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

Lines changed: 127 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -31,13 +31,15 @@
3131
package com.google.auth.oauth2;
3232

3333
import static org.junit.Assert.assertArrayEquals;
34+
import static org.junit.Assert.assertEquals;
3435
import static org.junit.Assert.assertThrows;
3536
import static org.junit.Assert.assertTrue;
3637

3738
import com.google.api.client.http.HttpStatusCodes;
3839
import com.google.auth.ServiceAccountSigner;
3940
import com.google.common.collect.ImmutableMap;
4041
import java.io.IOException;
42+
import org.junit.Before;
4143
import org.junit.Test;
4244
import org.junit.runner.RunWith;
4345
import org.junit.runners.JUnit4;
@@ -49,17 +51,53 @@ public class IamUtilsTest {
4951
private static final String CLIENT_EMAIL =
5052
"36680232662-vrd7ji19qe3nelgchd0ah2csanun6bnr@developer.gserviceaccount.com";
5153

54+
private ServiceAccountCredentials credentials;
55+
56+
@Before
57+
public void setup() throws IOException {
58+
// Mock this call for the Credentials because the IAM SignBlob RPC requires an access token. The
59+
// call is initialized with HttpCredentialsAdapter which will make a call to get the access
60+
// token
61+
credentials = Mockito.mock(ServiceAccountCredentials.class);
62+
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
63+
}
64+
5265
@Test
53-
public void sign_noRetry() throws IOException {
66+
public void sign_success_noRetry() {
5467
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
5568

56-
// Mock this call because signing requires an access token. The call is initialized with
57-
// HttpCredentialsAdapter which will make a call to get the access token
58-
ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
59-
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
69+
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
70+
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
71+
transportFactory.transport.setSignedBlob(expectedSignature);
72+
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
73+
transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
74+
75+
byte[] signature =
76+
IamUtils.sign(
77+
CLIENT_EMAIL,
78+
credentials,
79+
transportFactory.transport,
80+
expectedSignature,
81+
ImmutableMap.of());
82+
assertArrayEquals(expectedSignature, signature);
83+
84+
assertEquals(1, transportFactory.transport.getNumRequests());
85+
}
86+
87+
// The SignBlob RPC will retry up to three times before it gives up. This test will return two
88+
// 5xx status codes before returning a success. This test covers the cases where the number of
89+
// retry attempts is below the configured retry attempt count bounds (3 attempts).
90+
@Test
91+
public void sign_retryTwoTimes_success() {
92+
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
6093

6194
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
6295
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
96+
transportFactory.transport.addStatusCodeAndMessage(
97+
HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
98+
transportFactory.transport.addStatusCodeAndMessage(
99+
HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE, "Unavailable");
100+
transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
63101
transportFactory.transport.setSignedBlob(expectedSignature);
64102
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
65103

@@ -71,22 +109,97 @@ public void sign_noRetry() throws IOException {
71109
expectedSignature,
72110
ImmutableMap.of());
73111
assertArrayEquals(expectedSignature, signature);
112+
113+
// Expect that three requests are made (2 failures which are retries + 1 final requests which
114+
// resulted in a successful response)
115+
assertEquals(3, transportFactory.transport.getNumRequests());
74116
}
75117

118+
// The rpc will retry up to three times before it gives up. This test will enqueue three failed
119+
// status codes + messages before returning a success. After the third retry attempt, the request
120+
// will try one last time and the result will be reported back to the user.
76121
@Test
77-
public void sign_4xxServerError_exception() throws IOException {
122+
public void sign_retryThreeTimes_success() {
78123
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
79124

80-
// Mock this call because signing requires an access token. The call is initialized with
81-
// HttpCredentialsAdapter which will make a call to get the access token
82-
ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
83-
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
125+
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
126+
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
127+
transportFactory.transport.setSignedBlob(expectedSignature);
128+
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
129+
transportFactory.transport.addStatusCodeAndMessage(
130+
HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
131+
transportFactory.transport.addStatusCodeAndMessage(
132+
HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE, "Unavailable");
133+
transportFactory.transport.addStatusCodeAndMessage(
134+
HttpStatusCodes.STATUS_CODE_SERVER_ERROR, "Server Error");
135+
transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
136+
137+
byte[] signature =
138+
IamUtils.sign(
139+
CLIENT_EMAIL,
140+
credentials,
141+
transportFactory.transport,
142+
expectedSignature,
143+
ImmutableMap.of());
144+
assertArrayEquals(expectedSignature, signature);
145+
146+
// Expect that three requests are made (3 failures which are retried + 1 final request which
147+
// resulted the final success response)
148+
assertEquals(4, transportFactory.transport.getNumRequests());
149+
}
150+
151+
// The rpc will retry up to three times before it gives up. This test will enqueue four failed
152+
// status codes + messages before returning a success. After the third retry attempt, the request
153+
// will try one last time and the result will be reported back to the user.
154+
@Test
155+
public void sign_retryThreeTimes_exception() {
156+
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
84157

85158
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
86159
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
87160
transportFactory.transport.setSignedBlob(expectedSignature);
88161
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
89-
transportFactory.transport.setErrorResponseCodeAndMessage(
162+
transportFactory.transport.addStatusCodeAndMessage(
163+
HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
164+
transportFactory.transport.addStatusCodeAndMessage(
165+
HttpStatusCodes.STATUS_CODE_SERVICE_UNAVAILABLE, "Unavailable");
166+
transportFactory.transport.addStatusCodeAndMessage(
167+
HttpStatusCodes.STATUS_CODE_SERVER_ERROR, "Server Error");
168+
transportFactory.transport.addStatusCodeAndMessage(
169+
HttpStatusCodes.STATUS_CODE_BAD_GATEWAY, "Bad Gateway");
170+
transportFactory.transport.addStatusCodeAndMessage(HttpStatusCodes.STATUS_CODE_OK, "");
171+
172+
ServiceAccountSigner.SigningException exception =
173+
assertThrows(
174+
ServiceAccountSigner.SigningException.class,
175+
() ->
176+
IamUtils.sign(
177+
CLIENT_EMAIL,
178+
credentials,
179+
transportFactory.transport,
180+
expectedSignature,
181+
ImmutableMap.of()));
182+
assertTrue(exception.getMessage().contains("Failed to sign the provided bytes"));
183+
assertTrue(
184+
exception
185+
.getCause()
186+
.getMessage()
187+
.contains("Unexpected Error code 502 trying to sign provided bytes"));
188+
189+
// Expect that three requests are made (3 failures which are retried + 1 final request which
190+
// resulted in another failed response)
191+
assertEquals(4, transportFactory.transport.getNumRequests());
192+
}
193+
194+
@Test
195+
public void sign_4xxError_noRetry_exception() {
196+
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
197+
198+
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
199+
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
200+
transportFactory.transport.setSignedBlob(expectedSignature);
201+
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
202+
transportFactory.transport.addStatusCodeAndMessage(
90203
HttpStatusCodes.STATUS_CODE_UNAUTHORIZED, "Failed to sign the provided bytes");
91204

92205
ServiceAccountSigner.SigningException exception =
@@ -105,5 +218,8 @@ public void sign_4xxServerError_exception() throws IOException {
105218
.getCause()
106219
.getMessage()
107220
.contains("Error code 401 trying to sign provided bytes:"));
221+
222+
// Only one request will have been made for a 4xx error (no retries)
223+
assertEquals(1, transportFactory.transport.getNumRequests());
108224
}
109225
}

0 commit comments

Comments
 (0)