Fix DefaultTokenRefresher.

Author: rosalyntan

appcheck/firebase-appcheck/src/main/java/com/google/firebase/appcheck/internal/DefaultFirebaseAppCheck.java

@@ -184,12 +184,8 @@ Task<AppCheckTokenResult> fetchTokenResultFromProvider() {
               @Override
               public Task<AppCheckTokenResult> then(@NonNull Task<AppCheckToken> task) {
                 if (task.isSuccessful()) {
-                  AppCheckTokenResult tokenResult =
-                      DefaultAppCheckTokenResult.constructFromAppCheckToken(task.getResult());
-                  for (AppCheckTokenListener listener : appCheckTokenListenerList) {
-                    listener.onAppCheckTokenChanged(tokenResult);
-                  }
-                  return Tasks.forResult(tokenResult);
+                  return Tasks.forResult(
+                      DefaultAppCheckTokenResult.constructFromAppCheckToken(task.getResult()));
                 }
                 // If the token exchange failed, return a dummy token for integrators to attach in
                 // their headers.
@@ -215,6 +211,11 @@ public Task<AppCheckToken> then(@NonNull Task<AppCheckToken> task) {
                   for (AppCheckListener listener : appCheckListenerList) {
                     listener.onAppCheckTokenChanged(token);
                   }
+                  AppCheckTokenResult tokenResult =
+                      DefaultAppCheckTokenResult.constructFromAppCheckToken(token);
+                  for (AppCheckTokenListener listener : appCheckTokenListenerList) {
+                    listener.onAppCheckTokenChanged(tokenResult);
+                  }
                 }
                 return task;
               }

appcheck/firebase-appcheck/src/main/java/com/google/firebase/appcheck/internal/DefaultTokenRefresher.java

@@ -20,10 +20,9 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
-import com.google.android.gms.tasks.OnCompleteListener;
+import com.google.android.gms.tasks.OnFailureListener;
 import com.google.android.gms.tasks.Task;
-import com.google.firebase.appcheck.AppCheckTokenResult;
-import com.google.firebase.appcheck.internal.util.Logger;
+import com.google.firebase.appcheck.AppCheckToken;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -94,20 +93,12 @@ private long getNextRefreshMillis() {
   }
 
   private void onRefresh() {
-    Task<AppCheckTokenResult> task = firebaseAppCheck.fetchTokenFromProvider();
-    task.addOnCompleteListener(
-        new OnCompleteListener<AppCheckTokenResult>() {
+    Task<AppCheckToken> task = firebaseAppCheck.fetchTokenFromProvider();
+    task.addOnFailureListener(
+        new OnFailureListener() {
           @Override
-          public void onComplete(@NonNull Task<AppCheckTokenResult> task) {
-            if (task.isSuccessful()) {
-              AppCheckTokenResult tokenResult = task.getResult();
-              if (tokenResult.getError() != null) {
-                scheduleRefreshAfterFailure();
-              }
-            } else {
-              // Task was not successful; this should not happen.
-              Logger.getLogger().e("Unexpected failure while fetching token.");
-            }
+          public void onFailure(@NonNull Exception e) {
+            scheduleRefreshAfterFailure();
           }
         });
   }

appcheck/firebase-appcheck/src/test/java/com/google/firebase/appcheck/internal/DefaultTokenRefresherTest.java

@@ -24,8 +24,7 @@
 import static org.mockito.Mockito.when;
 
 import com.google.android.gms.tasks.Tasks;
-import com.google.firebase.FirebaseException;
-import com.google.firebase.appcheck.AppCheckTokenResult;
+import com.google.firebase.appcheck.AppCheckToken;
 import java.util.concurrent.ScheduledExecutorService;
 import org.junit.Before;
 import org.junit.Test;
@@ -47,7 +46,7 @@ public class DefaultTokenRefresherTest {
 
   @Mock DefaultFirebaseAppCheck mockFirebaseAppCheck;
   @Mock ScheduledExecutorService mockScheduledExecutorService;
-  @Mock AppCheckTokenResult mockAppCheckTokenResult;
+  @Mock AppCheckToken mockAppCheckToken;
 
   private DefaultTokenRefresher defaultTokenRefresher;
 
@@ -56,7 +55,7 @@ public void setUp() {
     MockitoAnnotations.initMocks(this);
 
     when(mockFirebaseAppCheck.fetchTokenFromProvider())
-        .thenReturn(Tasks.forResult(mockAppCheckTokenResult));
+        .thenReturn(Tasks.forResult(mockAppCheckToken));
 
     defaultTokenRefresher =
         new DefaultTokenRefresher(mockFirebaseAppCheck, mockScheduledExecutorService);
@@ -77,7 +76,8 @@ public void scheduleRefresh_success() {
 
   @Test
   public void scheduleRefresh_taskFails_schedulesRefreshAfterFailure() {
-    when(mockAppCheckTokenResult.getError()).thenReturn(new FirebaseException(ERROR));
+    when(mockFirebaseAppCheck.fetchTokenFromProvider())
+        .thenReturn(Tasks.forException(new Exception()));
     defaultTokenRefresher.scheduleRefresh(TIME_TO_REFRESH_MILLIS);
 
     ArgumentCaptor<Runnable> onRefreshCaptor = ArgumentCaptor.forClass(Runnable.class);
@@ -94,7 +94,8 @@ public void scheduleRefresh_taskFails_schedulesRefreshAfterFailure() {
 
   @Test
   public void scheduleRefreshAfterFailure_exponentialBackoff() {
-    when(mockAppCheckTokenResult.getError()).thenReturn(new FirebaseException(ERROR));
+    when(mockFirebaseAppCheck.fetchTokenFromProvider())
+        .thenReturn(Tasks.forException(new Exception()));
     defaultTokenRefresher.scheduleRefresh(TIME_TO_REFRESH_MILLIS);
 
     ArgumentCaptor<Runnable> onRefreshCaptor = ArgumentCaptor.forClass(Runnable.class);
@@ -131,7 +132,8 @@ public void scheduleRefreshAfterFailure_exponentialBackoff() {
 
   @Test
   public void scheduleRefresh_resetsDelay() {
-    when(mockAppCheckTokenResult.getError()).thenReturn(new FirebaseException(ERROR));
+    when(mockFirebaseAppCheck.fetchTokenFromProvider())
+        .thenReturn(Tasks.forException(new Exception()));
     defaultTokenRefresher.scheduleRefresh(TIME_TO_REFRESH_MILLIS);
 
     ArgumentCaptor<Runnable> onRefreshCaptor = ArgumentCaptor.forClass(Runnable.class);