From a6c0a2675fea56c80b208de20e43d5121eae9b93 Mon Sep 17 00:00:00 2001
From: Di Wu <github@wudi.me>
Date: Fri, 12 Apr 2024 15:13:31 -0700
Subject: [PATCH] fix(datastore): flacky unit test

---
 .../Sync/Support/MutationEventQueryTests.swift            | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Sync/Support/MutationEventQueryTests.swift b/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Sync/Support/MutationEventQueryTests.swift
index 9e1eb5df65..e982872609 100644
--- a/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Sync/Support/MutationEventQueryTests.swift
+++ b/AmplifyPlugins/DataStore/Tests/AWSDataStorePluginTests/Sync/Support/MutationEventQueryTests.swift
@@ -58,7 +58,7 @@ class MutationEventQueryTests: BaseDataStoreTests {
         wait(for: [querySuccess], timeout: 1)
     }
 
-    func testQueryPendingMutationEventsForModelIds() {
+    func testQueryPendingMutationEventsForModelIds() async {
         let mutationEvent1 = generateRandomMutationEvent()
         let mutationEvent2 = generateRandomMutationEvent()
 
@@ -70,7 +70,7 @@ class MutationEventQueryTests: BaseDataStoreTests {
             }
             saveMutationEvent1.fulfill()
         }
-        wait(for: [saveMutationEvent1], timeout: 1)
+        await fulfillment(of: [saveMutationEvent1], timeout: 1)
 
         let saveMutationEvent2 = expectation(description: "save mutationEvent1 success")
         storageAdapter.save(mutationEvent2) { result in
@@ -80,7 +80,7 @@ class MutationEventQueryTests: BaseDataStoreTests {
             }
             saveMutationEvent2.fulfill()
         }
-        wait(for: [saveMutationEvent2], timeout: 1)
+        await fulfillment(of: [saveMutationEvent2], timeout: 1)
 
         let querySuccess = expectation(description: "query for metadata success")
         var mutationEvents = [mutationEvent1]
@@ -98,7 +98,7 @@ class MutationEventQueryTests: BaseDataStoreTests {
             }
         }
 
-        wait(for: [querySuccess], timeout: 1)
+        await fulfillment(of: [querySuccess], timeout: 5)
     }
 
     private func generateRandomMutationEvent() -> MutationEvent {