Enabled automatic retries for FirebaseUserManager (#355)

Author: hiranya911

src/main/java/com/google/firebase/auth/FirebaseAuth.java

@@ -75,12 +75,7 @@ private FirebaseAuth(Builder builder) {
     this.tokenFactory = threadSafeMemoize(builder.tokenFactory);
     this.idTokenVerifier = threadSafeMemoize(builder.idTokenVerifier);
     this.cookieVerifier = threadSafeMemoize(builder.cookieVerifier);
-    this.userManager = threadSafeMemoize(new Supplier<FirebaseUserManager>() {
-      @Override
-      public FirebaseUserManager get() {
-        return new FirebaseUserManager(firebaseApp);
-      }
-    });
+    this.userManager = threadSafeMemoize(builder.userManager);
     this.jsonFactory = firebaseApp.getOptions().getJsonFactory();
   }
 
@@ -1109,10 +1104,10 @@ protected String execute() throws FirebaseAuthException {
   }
 
   private <T> Supplier<T> threadSafeMemoize(final Supplier<T> supplier) {
-    checkNotNull(supplier);
     return Suppliers.memoize(new Supplier<T>() {
       @Override
       public T get() {
+        checkNotNull(supplier);
         synchronized (lock) {
           checkNotDestroyed();
           return supplier.get();
@@ -1155,6 +1150,12 @@ public FirebaseTokenVerifier get() {
             return FirebaseTokenUtils.createSessionCookieVerifier(app, Clock.SYSTEM);
           }
         })
+        .setUserManager(new Supplier<FirebaseUserManager>() {
+          @Override
+          public FirebaseUserManager get() {
+            return new FirebaseUserManager(app);
+          }
+        })
         .build();
   }
 
@@ -1168,6 +1169,7 @@ static class Builder {
     private Supplier<FirebaseTokenFactory> tokenFactory;
     private Supplier<? extends FirebaseTokenVerifier> idTokenVerifier;
     private Supplier<? extends FirebaseTokenVerifier> cookieVerifier;
+    private Supplier<FirebaseUserManager> userManager;
 
     private Builder() { }
 
@@ -1191,6 +1193,11 @@ Builder setCookieVerifier(Supplier<? extends FirebaseTokenVerifier> cookieVerifi
       return this;
     }
 
+    Builder setUserManager(Supplier<FirebaseUserManager> userManager) {
+      this.userManager = userManager;
+      return this;
+    }
+
     FirebaseAuth build() {
       return new FirebaseAuth(this);
     }

src/main/java/com/google/firebase/auth/FirebaseUserManager.java

@@ -26,7 +26,6 @@
 import com.google.api.client.http.HttpResponse;
 import com.google.api.client.http.HttpResponseException;
 import com.google.api.client.http.HttpResponseInterceptor;
-import com.google.api.client.http.HttpTransport;
 import com.google.api.client.http.json.JsonHttpContent;
 import com.google.api.client.json.GenericJson;
 import com.google.api.client.json.JsonFactory;
@@ -45,7 +44,7 @@
 
 import com.google.firebase.auth.internal.HttpErrorResponse;
 import com.google.firebase.auth.internal.UploadAccountResponse;
-import com.google.firebase.internal.FirebaseRequestInitializer;
+import com.google.firebase.internal.ApiClientUtils;
 import com.google.firebase.internal.NonNull;
 import com.google.firebase.internal.Nullable;
 import com.google.firebase.internal.SdkUtils;
@@ -111,6 +110,10 @@ class FirebaseUserManager {
    * @param app A non-null {@link FirebaseApp}.
    */
   FirebaseUserManager(@NonNull FirebaseApp app) {
+    this(app, null);
+  }
+
+  FirebaseUserManager(@NonNull FirebaseApp app, @Nullable HttpRequestFactory requestFactory) {
     checkNotNull(app, "FirebaseApp must not be null");
     String projectId = ImplFirebaseTrampolines.getProjectId(app);
     checkArgument(!Strings.isNullOrEmpty(projectId),
@@ -119,8 +122,12 @@ class FirebaseUserManager {
             + "set the project ID via the GOOGLE_CLOUD_PROJECT environment variable.");
     this.baseUrl = String.format(ID_TOOLKIT_URL, projectId);
     this.jsonFactory = app.getOptions().getJsonFactory();
-    HttpTransport transport = app.getOptions().getHttpTransport();
-    this.requestFactory = transport.createRequestFactory(new FirebaseRequestInitializer(app));
+
+    if (requestFactory == null) {
+      requestFactory = ApiClientUtils.newAuthorizedRequestFactory(app);
+    }
+
+    this.requestFactory = requestFactory;
   }
 
   @VisibleForTesting

src/test/java/com/google/firebase/auth/FirebaseAuthTest.java

@@ -429,11 +429,11 @@ private FirebaseAuth getAuthForIdTokenVerification(FirebaseTokenVerifier tokenVe
   private FirebaseAuth getAuthForIdTokenVerification(
       Supplier<? extends FirebaseTokenVerifier> tokenVerifierSupplier) {
     FirebaseApp app = FirebaseApp.initializeApp(firebaseOptions);
+    FirebaseUserManager userManager = new FirebaseUserManager(app);
     return FirebaseAuth.builder()
         .setFirebaseApp(app)
-        .setTokenFactory(Suppliers.<FirebaseTokenFactory>ofInstance(null))
         .setIdTokenVerifier(tokenVerifierSupplier)
-        .setCookieVerifier(Suppliers.<FirebaseTokenVerifier>ofInstance(null))
+        .setUserManager(Suppliers.ofInstance(userManager))
         .build();
   }
 
@@ -444,11 +444,11 @@ private FirebaseAuth getAuthForSessionCookieVerification(FirebaseTokenVerifier t
   private FirebaseAuth getAuthForSessionCookieVerification(
       Supplier<? extends FirebaseTokenVerifier> tokenVerifierSupplier) {
     FirebaseApp app = FirebaseApp.initializeApp(firebaseOptions);
+    FirebaseUserManager userManager = new FirebaseUserManager(app);
     return FirebaseAuth.builder()
         .setFirebaseApp(app)
-        .setTokenFactory(Suppliers.<FirebaseTokenFactory>ofInstance(null))
-        .setIdTokenVerifier(Suppliers.<FirebaseTokenVerifier>ofInstance(null))
         .setCookieVerifier(tokenVerifierSupplier)
+        .setUserManager(Suppliers.ofInstance(userManager))
         .build();
   }
 

src/test/java/com/google/firebase/auth/FirebaseUserManagerTest.java

@@ -33,6 +33,7 @@
 import com.google.api.client.testing.http.MockHttpTransport;
 import com.google.api.client.testing.http.MockLowLevelHttpResponse;
 import com.google.auth.oauth2.GoogleCredentials;
+import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
@@ -550,14 +551,7 @@ public void call(FirebaseAuth auth) throws Exception {
         .build();
 
     MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
-    MockHttpTransport transport = new MockHttpTransport.Builder()
-        .setLowLevelHttpResponse(response)
-        .build();
-    FirebaseApp.initializeApp(new FirebaseOptions.Builder()
-        .setCredentials(credentials)
-        .setProjectId("test-project-id")
-        .setHttpTransport(transport)
-        .build());
+    FirebaseAuth auth = getRetryDisabledAuth(response);
 
     // Test for common HTTP error codes
     for (int code : ImmutableList.of(302, 400, 401, 404, 500)) {
@@ -566,7 +560,7 @@ public void call(FirebaseAuth auth) throws Exception {
         response.setContent("{}");
         response.setStatusCode(code);
         try {
-          operation.call(FirebaseAuth.getInstance());
+          operation.call(auth);
           fail("No error thrown for HTTP error: " + code);
         } catch (ExecutionException e) {
           assertTrue(e.getCause() instanceof FirebaseAuthException);
@@ -584,7 +578,7 @@ public void call(FirebaseAuth auth) throws Exception {
       response.setContent("{\"error\": {\"message\": \"USER_NOT_FOUND\"}}");
       response.setStatusCode(500);
       try {
-        operation.call(FirebaseAuth.getInstance());
+        operation.call(auth);
         fail("No error thrown for HTTP error");
       }  catch (ExecutionException e) {
         assertTrue(e.getCause().toString(), e.getCause() instanceof FirebaseAuthException);
@@ -615,16 +609,9 @@ public void testGetUserUnexpectedHttpError() throws Exception {
     MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
     response.setContent("{\"not\" json}");
     response.setStatusCode(500);
-    MockHttpTransport transport = new MockHttpTransport.Builder()
-        .setLowLevelHttpResponse(response)
-        .build();
-    FirebaseApp.initializeApp(new FirebaseOptions.Builder()
-        .setCredentials(credentials)
-        .setProjectId("test-project-id")
-        .setHttpTransport(transport)
-        .build());
+    FirebaseAuth auth = getRetryDisabledAuth(response);
     try {
-      FirebaseAuth.getInstance().getUserAsync("testuser").get();
+      auth.getUserAsync("testuser").get();
       fail("No error thrown for JSON error");
     }  catch (ExecutionException e) {
       assertTrue(e.getCause() instanceof FirebaseAuthException);
@@ -1173,15 +1160,10 @@ public void testGenerateSignInWithEmailLinkWithSettings() throws Exception {
 
   @Test
   public void testHttpErrorWithCode() {
-    FirebaseApp.initializeApp(new FirebaseOptions.Builder()
-        .setCredentials(credentials)
-        .setHttpTransport(new MultiRequestMockHttpTransport(ImmutableList.of(
-            new MockLowLevelHttpResponse()
-                .setContent("{\"error\": {\"message\": \"UNAUTHORIZED_DOMAIN\"}}")
-                .setStatusCode(500))))
-        .setProjectId("test-project-id")
-        .build());
-    FirebaseAuth auth = FirebaseAuth.getInstance();
+    MockLowLevelHttpResponse response = new MockLowLevelHttpResponse()
+        .setContent("{\"error\": {\"message\": \"UNAUTHORIZED_DOMAIN\"}}")
+        .setStatusCode(500);
+    FirebaseAuth auth = getRetryDisabledAuth(response);
     FirebaseUserManager userManager = auth.getUserManager();
     try {
       userManager.getEmailActionLink(EmailLinkType.PASSWORD_RESET, "test@example.com", null);
@@ -1194,15 +1176,10 @@ public void testHttpErrorWithCode() {
 
   @Test
   public void testUnexpectedHttpError() {
-    FirebaseApp.initializeApp(new FirebaseOptions.Builder()
-        .setCredentials(credentials)
-        .setHttpTransport(new MultiRequestMockHttpTransport(ImmutableList.of(
-            new MockLowLevelHttpResponse()
-                .setContent("{}")
-                .setStatusCode(500))))
-        .setProjectId("test-project-id")
-        .build());
-    FirebaseAuth auth = FirebaseAuth.getInstance();
+    MockLowLevelHttpResponse response = new MockLowLevelHttpResponse()
+        .setContent("{}")
+        .setStatusCode(500);
+    FirebaseAuth auth = getRetryDisabledAuth(response);
     FirebaseUserManager userManager = auth.getUserManager();
     try {
       userManager.getEmailActionLink(EmailLinkType.PASSWORD_RESET, "test@example.com", null);
@@ -1231,6 +1208,26 @@ private static TestResponseInterceptor initializeAppForUserManagement(String ...
     return interceptor;
   }
 
+  private static FirebaseAuth getRetryDisabledAuth(MockLowLevelHttpResponse response) {
+    final MockHttpTransport transport = new MockHttpTransport.Builder()
+        .setLowLevelHttpResponse(response)
+        .build();
+    final FirebaseApp app = FirebaseApp.initializeApp(new FirebaseOptions.Builder()
+        .setCredentials(credentials)
+        .setProjectId("test-project-id")
+        .setHttpTransport(transport)
+        .build());
+    return FirebaseAuth.builder()
+        .setFirebaseApp(app)
+        .setUserManager(new Supplier<FirebaseUserManager>() {
+          @Override
+          public FirebaseUserManager get() {
+            return new FirebaseUserManager(app, transport.createRequestFactory());
+          }
+        })
+        .build();
+  }
+
   private static void checkUserRecord(UserRecord userRecord) {
     assertEquals("testuser", userRecord.getUid());
     assertEquals("testuser@example.com", userRecord.getEmail());