Fixed DispatchGroup’s Excessive Wait in FunctionsContext

FirebaseFunctions/Sources/Internal/FunctionsContext.swift

@@ -66,11 +66,17 @@ struct FunctionsContextProvider {
       dispatchGroup.enter()
 
       if options?.requireLimitedUseAppCheckTokens == true {
-        appCheck.getLimitedUseToken? { tokenResult in
-          // Send only valid token to functions.
-          if tokenResult.error == nil {
-            limitedUseAppCheckToken = tokenResult.token
+        // `getLimitedUseToken(completion:)` is an optional protocol method.
+        // If it’s not implemented, we still need to leave the dispatch group.
+        if let limitedUseTokenClosure = appCheck.getLimitedUseToken {
+          limitedUseTokenClosure { tokenResult in
+            // Send only valid token to functions.
+            if tokenResult.error == nil {
+              limitedUseAppCheckToken = tokenResult.token
+            }
+            dispatchGroup.leave()
           }
+        } else {
           dispatchGroup.leave()
         }
       } else {

FirebaseFunctions/Tests/Unit/ContextProviderTests.swift

@@ -107,6 +107,25 @@ class ContextProviderTests: XCTestCase {
     waitForExpectations(timeout: 0.1)
   }
 
+  func testContextWithAppCheckWithoutOptionalMethods() {
+    let appCheck = AppCheckFakeWithoutOptionalMethods(tokenResult: appCheckTokenSuccess)
+    let provider = FunctionsContextProvider(auth: nil, messaging: nil, appCheck: appCheck)
+    let expectation =
+      expectation(description: "Verify non-implemented method for limited-use tokens")
+    provider.getContext(options: .init(requireLimitedUseAppCheckTokens: true)) { context, error in
+      XCTAssertNotNil(context)
+      XCTAssertNil(error)
+      XCTAssertNil(context.authToken)
+      XCTAssertNil(context.fcmToken)
+      XCTAssertNil(context.appCheckToken)
+      // If the method for limited-use tokens is not implemented, the value should be `nil`:
+      XCTAssertNil(context.limitedUseAppCheckToken)
+      expectation.fulfill()
+    }
+    // Importantly, `getContext(options:_:)` must still finish in a timely manner:
+    waitForExpectations(timeout: 0.1)
+  }
+
   func testAllContextsAvailableSuccess() {
     appCheckFake.tokenResult = appCheckTokenSuccess
     let auth = FIRAuthInteropFake(token: "token", userID: "userID", error: nil)
@@ -149,3 +168,21 @@ class ContextProviderTests: XCTestCase {
     waitForExpectations(timeout: 0.1)
   }
 }
+
+// MARK: - Utilities
+
+private class AppCheckFakeWithoutOptionalMethods: NSObject, AppCheckInterop {
+  let tokenResult: FIRAppCheckTokenResultInterop
+
+  init(tokenResult: FIRAppCheckTokenResultInterop) {
+    self.tokenResult = tokenResult
+  }
+
+  func getToken(forcingRefresh: Bool, completion handler: @escaping AppCheckTokenHandlerInterop) {
+    handler(tokenResult)
+  }
+
+  func tokenDidChangeNotificationName() -> String { "AppCheckFakeTokenDidChangeNotification" }
+  func notificationTokenKey() -> String { "AppCheckFakeTokenNotificationKey" }
+  func notificationAppNameKey() -> String { "AppCheckFakeAppNameNotificationKey" }
+}

SharedTestUtilities/AppCheckFake/FIRAppCheckFake.m

@@ -45,15 +45,15 @@ - (void)getLimitedUseTokenWithCompletion:(FIRAppCheckTokenHandlerInterop)handler
 }
 
 - (nonnull NSString *)notificationAppNameKey {
-  return @"FakeAppCheckTokenDidChangeNotification";
+  return @"AppCheckFakeAppNameNotificationKey";
 }
 
 - (nonnull NSString *)notificationTokenKey {
-  return @"FakeTokenNotificationKey";
+  return @"AppCheckFakeTokenNotificationKey";
 }
 
 - (nonnull NSString *)tokenDidChangeNotificationName {
-  return @"FakeAppCheckTokenDidChangeNotification";
+  return @"AppCheckFakeTokenDidChangeNotification";
 }
 
 @end