chore: Add tests for ClientCalls

googleapis · Jan 8, 2024 · 558e3fb · 558e3fb
1 parent 414483a
commit 558e3fb
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 9 deletions.
diff --git a/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java b/gax-java/gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcCallContext.java
@@ -630,7 +630,7 @@ public void validateUniverseDomain() {
         throw ApiExceptionFactory.createException(
             new Throwable("Unable to retrieve the Universe Domain from the Credentials."),
             GrpcStatusCode.of(Status.Code.UNAVAILABLE),
-            false);
+            ((Retryable) e).isRetryable());
       }
       // Return the exception back to the user
       throw new RuntimeException(e);

diff --git a/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java b/gax-java/gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcClientCallsTest.java
@@ -37,7 +37,9 @@
 import com.google.api.gax.grpc.testing.FakeServiceGrpc;
 import com.google.api.gax.rpc.EndpointContext;
 import com.google.api.gax.rpc.PermissionDeniedException;
+import com.google.api.gax.rpc.UnavailableException;
 import com.google.auth.Credentials;
+import com.google.auth.Retryable;
 import com.google.common.collect.ImmutableList;
 import com.google.common.truth.Truth;
 import com.google.type.Color;
@@ -63,6 +65,26 @@
 
 public class GrpcClientCallsTest {
 
+  // Auth Library's GoogleAuthException is package-private. Copy basic functionality for tests
+  private static class GoogleAuthException extends IOException implements Retryable {
+
+    private final boolean isRetryable;
+
+    private GoogleAuthException(boolean isRetryable) {
+      this.isRetryable = isRetryable;
+    }
+
+    @Override
+    public boolean isRetryable() {
+      return isRetryable;
+    }
+
+    @Override
+    public int getRetryCount() {
+      return 0;
+    }
+  }
+
   private GrpcCallContext defaultCallContext;
   private EndpointContext endpointContext;
   private Credentials credentials;
@@ -83,9 +105,6 @@ public void setUp() throws IOException {
   public void testAffinity() throws IOException {
     MethodDescriptor<Color, Money> descriptor = FakeServiceGrpc.METHOD_RECOGNIZE;
 
-    EndpointContext endpointContext = Mockito.mock(EndpointContext.class);
-    Mockito.when(endpointContext.hasValidUniverseDomain(Mockito.any())).thenReturn(true);
-
     @SuppressWarnings("unchecked")
     ClientCall<Color, Money> clientCall0 = Mockito.mock(ClientCall.class);
 
@@ -258,7 +277,8 @@ public void testTimeoutBeforeDeadline() throws IOException {
   }
 
   @Test
-  public void testValidUniverseDomain() {
+  public void testValidUniverseDomain() throws IOException {
+    Mockito.when(endpointContext.hasValidUniverseDomain(credentials)).thenReturn(true);
     GrpcCallContext context =
         GrpcCallContext.createDefault()
             .withChannel(mockChannel)
@@ -272,6 +292,7 @@ public void testValidUniverseDomain() {
     Mockito.verify(mockChannel, Mockito.times(1)).newCall(descriptor, callOptions);
   }
 
+  // This test is when the universe domain does not match
   @Test
   public void testInvalidUniverseDomain() throws IOException {
     Mockito.when(endpointContext.hasValidUniverseDomain(credentials)).thenReturn(false);
@@ -291,4 +312,26 @@ public void testInvalidUniverseDomain() throws IOException {
         .isEqualTo(GrpcStatusCode.Code.PERMISSION_DENIED);
     Mockito.verify(mockChannel, Mockito.never()).newCall(descriptor, callOptions);
   }
+
+  // This test is when the MDS is unable to return a valid universe domain
+  @Test
+  public void testUniverseDomainNotReady_shouldRetry() throws IOException {
+    Mockito.when(endpointContext.hasValidUniverseDomain(credentials))
+        .thenThrow(new GoogleAuthException(true));
+    GrpcCallContext context =
+        GrpcCallContext.createDefault()
+            .withChannel(mockChannel)
+            .withCredentials(credentials)
+            .withEndpointContext(endpointContext);
+
+    CallOptions callOptions = context.getCallOptions();
+
+    MethodDescriptor<Color, Money> descriptor = FakeServiceGrpc.METHOD_RECOGNIZE;
+    UnavailableException exception =
+        assertThrows(
+            UnavailableException.class, () -> GrpcClientCalls.newCall(descriptor, context));
+    assertThat(exception.getStatusCode().getCode()).isEqualTo(GrpcStatusCode.Code.UNAVAILABLE);
+    Truth.assertThat(exception.isRetryable()).isTrue();
+    Mockito.verify(mockChannel, Mockito.never()).newCall(descriptor, callOptions);
+  }
 }
diff --git a/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonClientCallsTest.java b/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonClientCallsTest.java
@@ -35,6 +35,7 @@
 import com.google.api.gax.rpc.EndpointContext;
 import com.google.api.gax.rpc.PermissionDeniedException;
 import com.google.auth.Credentials;
+import com.google.auth.Retryable;
 import java.io.IOException;
 import org.junit.Before;
 import org.junit.Test;
@@ -45,6 +46,26 @@
 @RunWith(JUnit4.class)
 public class HttpJsonClientCallsTest {
 
+  // Auth Library's GoogleAuthException is package-private. Copy basic functionality for tests
+  private static class GoogleAuthException extends IOException implements Retryable {
+
+    private final boolean isRetryable;
+
+    private GoogleAuthException(boolean isRetryable) {
+      this.isRetryable = isRetryable;
+    }
+
+    @Override
+    public boolean isRetryable() {
+      return isRetryable;
+    }
+
+    @Override
+    public int getRetryCount() {
+      return 0;
+    }
+  }
+
   private Credentials credentials;
   private EndpointContext endpointContext;
   private HttpJsonChannel mockChannel;
@@ -65,8 +86,8 @@ public void setUp() throws IOException {
             .withEndpointContext(endpointContext)
             .withChannel(mockChannel);
 
-    Mockito.when(callOptions.getCredentials()).thenReturn(credentials);
     Mockito.when(endpointContext.hasValidUniverseDomain(credentials)).thenReturn(true);
+    Mockito.when(callOptions.getCredentials()).thenReturn(credentials);
   }
 
   @Test
@@ -75,6 +96,7 @@ public void testValidUniverseDomain() {
     Mockito.verify(mockChannel, Mockito.times(1)).newCall(descriptor, callOptions);
   }
 
+  // This test is when the universe domain does not match
   @Test
   public void testInvalidUniverseDomain() throws IOException {
     Mockito.when(endpointContext.hasValidUniverseDomain(credentials)).thenReturn(false);
@@ -87,4 +109,22 @@ public void testInvalidUniverseDomain() throws IOException {
         .isEqualTo(HttpJsonStatusCode.Code.PERMISSION_DENIED);
     Mockito.verify(mockChannel, Mockito.never()).newCall(descriptor, callOptions);
   }
+
+  // This test is when the MDS is unable to return a valid universe domain
+  @Test
+  public void testUniverseDomainNotReady_shouldRetry() throws IOException {
+    Mockito.when(endpointContext.hasValidUniverseDomain(credentials))
+        .thenThrow(new GoogleAuthException(true));
+    HttpJsonCallContext context =
+        HttpJsonCallContext.createDefault()
+            .withChannel(mockChannel)
+            .withCredentials(credentials)
+            .withEndpointContext(endpointContext);
+
+    HttpJsonCallOptions callOptions = context.getCallOptions();
+
+    assertThrows(
+        RuntimeException.class, () -> HttpJsonClientCalls.newCall(descriptor, callContext));
+    Mockito.verify(mockChannel, Mockito.never()).newCall(descriptor, callOptions);
+  }
 }
diff --git a/gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java b/gax-java/gax/src/test/java/com/google/api/gax/rpc/EndpointContextTest.java
@@ -405,7 +405,8 @@ public void merge_nullInput() throws IOException {
 
   @Test
   public void merge_customEndpointContextInput() throws IOException {
-    EndpointContext originalEndpointContext = EndpointContext.newBuilder()
+    EndpointContext originalEndpointContext =
+        EndpointContext.newBuilder()
             .setServiceName("serviceName1")
             .setUniverseDomain("universeDomain1")
             .setClientSettingsEndpoint("clientSettingsEndpoint1")
@@ -414,7 +415,8 @@ public void merge_customEndpointContextInput() throws IOException {
             .setSwitchToMtlsEndpointAllowed(false)
             .setUsingGDCH(false)
             .build();
-    EndpointContext newEndpointContext = EndpointContext.newBuilder()
+    EndpointContext newEndpointContext =
+        EndpointContext.newBuilder()
             .setServiceName("serviceName2")
             .setUniverseDomain("universeDomain2")
             .setClientSettingsEndpoint("clientSettingsEndpoint2")
@@ -424,6 +426,7 @@ public void merge_customEndpointContextInput() throws IOException {
             // Cannot set GDC-H as true as GDC-H is incompatible with Universe Domain
             .setUsingGDCH(false)
             .build();
-    Truth.assertThat(originalEndpointContext.merge(newEndpointContext)).isEqualTo(newEndpointContext);
+    Truth.assertThat(originalEndpointContext.merge(newEndpointContext))
+        .isEqualTo(newEndpointContext);
   }
 }