fix(fdc): Don't throw when FirebaseAuth is unable to get an ID Token (#16711)

maneesht · web-flow · commit 1ef2044a7a9f · 2024-11-19T08:53:15.000-08:00
diff --git a/packages/firebase_data_connect/firebase_data_connect/lib/src/core/ref.dart b/packages/firebase_data_connect/firebase_data_connect/lib/src/core/ref.dart
@@ -48,8 +48,13 @@ abstract class OperationRef<Data, Variables> {
 
   Future<OperationResult<Data, Variables>> execute();
   Future<bool> _shouldRetry() async {
-    String? newToken =
-        await this.dataConnect.auth?.currentUser?.getIdToken(false);
+    String? newToken;
+    try {
+      newToken = await this.dataConnect.auth?.currentUser?.getIdToken(false);
+    } catch (e) {
+      // Don't retry if there was an issue getting the ID Token.
+      log("There was an error attempting to retrieve the ID Token: ${e.toString()}");
+    }
     bool shouldRetry = newToken != null && _lastToken != newToken;
     _lastToken = newToken;
     return shouldRetry;
@@ -97,12 +102,14 @@ class QueryManager {
       return;
     }
     StreamController stream = trackedQueries[operationName]![key]!;
-    // TODO(mtewani): Prevent this from getting called multiple times.
-    stream.onCancel = () => stream.close();
-    if (error != null) {
-      stream.addError(error);
-    } else {
-      stream.add(QueryResult<Data, Variables>(dataConnect, data as Data, ref));
+
+    if (!stream.isClosed) {
+      if (error != null) {
+        stream.addError(error);
+      } else {
+        stream
+            .add(QueryResult<Data, Variables>(dataConnect, data as Data, ref));
+      }
     }
   }
 }
diff --git a/packages/firebase_data_connect/firebase_data_connect/test/src/core/ref_test.dart b/packages/firebase_data_connect/firebase_data_connect/test/src/core/ref_test.dart
@@ -32,20 +32,7 @@ class MockDataConnectTransport extends Mock implements DataConnectTransport {}
 
 class MockFirebaseDataConnect extends Mock implements FirebaseDataConnect {}
 
-class DCMockUser extends Mock implements User {
-  int count = 0;
-  List<String?> tokens = ['invalid-token', 'valid-token'];
-  @override
-  Future<String?> getIdToken([bool forceRefresh = false]) {
-    // First return an invalid token, then return a valid token
-    return Future.value(tokens[count++]);
-  }
-}
-
-class MockFirebaseAuth extends Mock implements FirebaseAuth {
-  @override
-  User? get currentUser => DCMockUser();
-}
+class MockFirebaseAuth extends Mock implements FirebaseAuth {}
 
 class MockQueryManager extends Mock implements QueryManager {}
 
@@ -122,10 +109,15 @@ void main() {
     late Serializer<String> serializer;
     late MockClient mockHttpClient;
     late Deserializer<String> deserializer;
+    late MockFirebaseAuth auth;
+    late MockUser mockUser;
 
     setUp(() {
       mockDataConnect = MockFirebaseDataConnect();
-      when(mockDataConnect.auth).thenReturn(MockFirebaseAuth());
+      auth = MockFirebaseAuth();
+      mockUser = MockUser();
+      when(mockDataConnect.auth).thenReturn(auth);
+      when(auth.currentUser).thenReturn(mockUser);
       mockHttpClient = MockClient();
       transport = RestTransport(
         TransportOptions('testhost', 443, true),
@@ -142,15 +134,31 @@ void main() {
       transport.setHttp(mockHttpClient);
       mockDataConnect.transport = transport;
     });
+    test('executeQuery should gracefully handle getIdToken failures', () async {
+      final deserializer = (String data) => 'Deserialized Data';
+      final mockResponseSuccess = http.Response('{"success": true}', 200);
+      when(mockUser.getIdToken()).thenThrow(Exception('Auth error'));
+      QueryRef ref = QueryRef(mockDataConnect, 'operation', transport,
+          deserializer, QueryManager(mockDataConnect), emptySerializer, null);
+      when(mockHttpClient.post(any,
+              headers: anyNamed('headers'), body: anyNamed('body')))
+          .thenAnswer((_) async => mockResponseSuccess);
+      await ref.execute();
+    });
     test(
         'query should forceRefresh on ID token if the first request is unauthorized',
         () async {
       final mockResponse = http.Response('{"error": "Unauthorized"}', 401);
       final mockResponseSuccess = http.Response('{"success": true}', 200);
       final deserializer = (String data) => 'Deserialized Data';
       int count = 0;
+      int idTokenCount = 0;
       QueryRef ref = QueryRef(mockDataConnect, 'operation', transport,
           deserializer, QueryManager(mockDataConnect), emptySerializer, null);
+      when(mockUser.getIdToken()).thenAnswer((invocation) => [
+            Future.value('invalid-token'),
+            Future.value('valid-token')
+          ][idTokenCount++]);
 
       when(mockHttpClient.post(any,
               headers: anyNamed('headers'), body: anyNamed('body')))
