[Infra] Reduce test flakes due to async variable assignment (#13610)

firebase · Sep 9, 2024 · b2a33ce · b2a33ce
1 parent 3d035c1
commit b2a33ce
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 15 deletions.
diff --git a/FirebaseAuth/Tests/Unit/AuthBackendRPCImplentationTests.swift b/FirebaseAuth/Tests/Unit/AuthBackendRPCImplentationTests.swift
@@ -590,12 +590,10 @@ class AuthBackendRPCImplementationTests: RPCBaseTests {
         try self.rpcIssuer.respond(withJSON: [:])
       }
       _ = try? await rpcImplementation.call(with: request)
-      // Make sure completeRequest updates.
-      usleep(10000)
 
       // Then
       let expectedHeader = HeartbeatLoggingTestUtils.nonEmptyHeartbeatsPayload.headerValue()
-      let completeRequest = try XCTUnwrap(rpcIssuer.completeRequest)
+      let completeRequest = await rpcIssuer.completeRequest.value
       let headerValue = completeRequest.value(forHTTPHeaderField: "X-Firebase-Client")
       XCTAssertEqual(headerValue, expectedHeader)
     }
@@ -619,10 +617,8 @@ class AuthBackendRPCImplementationTests: RPCBaseTests {
         try self.rpcIssuer.respond(withJSON: [:])
       }
       _ = try? await rpcImplementation.call(with: request)
-      // Make sure completeRequest updates.
-      usleep(10000)
 
-      let completeRequest = try XCTUnwrap(rpcIssuer.completeRequest)
+      let completeRequest = await rpcIssuer.completeRequest.value
       let headerValue = completeRequest.value(forHTTPHeaderField: "X-Firebase-AppCheck")
       XCTAssertEqual(headerValue, fakeAppCheck.fakeAppCheckToken)
     }
@@ -651,11 +647,9 @@ class AuthBackendRPCImplementationTests: RPCBaseTests {
         try self.rpcIssuer.respond(withJSON: [:])
       }
       _ = try? await rpcImplementation.call(with: request)
-      // Make sure completRequest updates.
-      usleep(10000)
 
       // Then
-      let completeRequest = try XCTUnwrap(rpcIssuer.completeRequest)
+      let completeRequest = await rpcIssuer.completeRequest.value
       XCTAssertNil(completeRequest.value(forHTTPHeaderField: "X-Firebase-Client"))
     }
   #endif // COCOAPODS || SWIFT_PACKAGE

diff --git a/FirebaseAuth/Tests/Unit/Fakes/FakeBackendRPCIssuer.swift b/FirebaseAuth/Tests/Unit/Fakes/FakeBackendRPCIssuer.swift
@@ -52,7 +52,7 @@ class FakeBackendRPCIssuer: NSObject, AuthBackendRPCIssuer {
   /** @property completeRequest
       @brief The last request to be processed by the backend.
    */
-  var completeRequest: URLRequest?
+  var completeRequest: Task<URLRequest, Never>!
 
   /** @var _handler
       @brief A block we must invoke when @c respondWithError or @c respondWithJSON are called.
@@ -148,11 +148,10 @@ class FakeBackendRPCIssuer: NSObject, AuthBackendRPCIssuer {
       requestData = body
       // Use the real implementation so that the complete request can
       // be verified during testing.
-      Task {
-        self.completeRequest = await AuthBackend.request(withURL: requestURL!,
-                                                         contentType: contentType,
-                                                         requestConfiguration: request
-                                                           .requestConfiguration())
+      completeRequest = Task {
+        await AuthBackend.request(withURL: requestURL!,
+                                  contentType: contentType,
+                                  requestConfiguration: request.requestConfiguration())
       }
       decodedRequest = try? JSONSerialization.jsonObject(with: body) as? [String: Any]
     }