test(datastore): fix datastore integration test (#2790)

Author: Di Wu

AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/AWSDataStorePlugin+DataStoreBaseBehavior.swift

@@ -500,7 +500,6 @@ extension AWSDataStorePlugin: DataStoreBaseBehavior {
     private func publishMutationEvent<M: Model>(from model: M,
                                                 modelSchema: ModelSchema,
                                                 mutationType: MutationEvent.MutationType) {
-
         let metadata = MutationSyncMetadata.keys
         let metadataId = MutationSyncMetadata.identifier(modelName: modelSchema.name,
                                                          modelId: model.identifier(schema: modelSchema).stringValue)

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/AWSDataStorePluginConfigurationTests.swift

@@ -6,11 +6,16 @@
 //
 
 import XCTest
-import Amplify
+@testable import Amplify
 import AWSDataStorePlugin
 
 class AWSDataStorePluginConfigurationTests: XCTestCase {
 
+    override func tearDown() async throws {
+        await Amplify.reset()
+        try await Task.sleep(seconds: 1)
+    }
+
     // Note this test requires the ability to write a new database in the Documents directcory, so it must be embedded
     // in a host app
     func testDoesNotThrowOnMissingConfig() throws {

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/Connection/DataStoreConnectionScenario6Tests.swift

@@ -205,7 +205,7 @@ class DataStoreConnectionScenario6Tests: SyncEngineIntegrationTestBase {
         let post2 = try await savePost(title: "title", blog: blog2)
         _ = try await saveComment(id: commentId2, post: post2, content: "content")
 
-        await waitForExpectations(timeout: 5)
+        await waitForExpectations(timeout: 10)
 
         let outboxMutationProcessed = expectation(description: "received outboxMutationProcessed")
         var processedSoFar = 0

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/DataStoreConfigurationTests.swift

@@ -14,13 +14,10 @@ import AWSPluginsCore
 
 class DataStoreConfigurationTests: XCTestCase {
 
-    override func setUp() async throws {
-        await Amplify.reset()
-    }
-
     override func tearDown() async throws {
         try await Amplify.DataStore.clear()
         await Amplify.reset()
+        try await Task.sleep(seconds: 1)
     }
 
     func testConfigureWithSameSchemaDoesNotDeleteDatabase() async throws {

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/DataStoreEndToEndTests.swift

@@ -489,50 +489,52 @@ class DataStoreEndToEndTests: SyncEngineIntegrationTestBase {
     func testParallelMutations_whenWaitingForEventToProcess_noMutationLoss() async throws {
         await setUp(withModels: TestModelRegistration())
         try await startAmplifyAndWaitForSync()
+
         let parallelSize = 100
-        let initExpectation = expectation(description: "expect MutationEventPublisher works fine")
-        let parallelExpectation = expectation(description: "expect parallel processing no data loss")
+        let initExpectation = asyncExpectation(description: "expect MutationEventPublisher works fine")
+        let parallelExpectation = asyncExpectation(description: "expect parallel processing no data loss")
 
         let newPost = Post(title: UUID().uuidString, content: UUID().uuidString, createdAt: .now())
 
         let titlePrefix = UUID().uuidString
         let posts = (0..<parallelSize).map { Post(title: "\(titlePrefix)-\($0)", content: UUID().uuidString, createdAt: .now()) }
         var expectedResult = Set<String>()
+        let extractPost: (DataStoreHubEvent) -> Post? = {
+            if case .outboxMutationProcessed(let mutationEvent) = $0,
+               mutationEvent.modelName == Post.modelName
+            {
+                return mutationEvent.element.model as? Post
+            }
+            return nil
+        }
+
         let cancellable = Amplify.Hub.publisher(for: .dataStore)
-            .sink { payload in
-                let event = DataStoreHubEvent(payload: payload)
-                switch event {
-                case .outboxMutationProcessed(let mutationEvent):
-                    guard mutationEvent.modelName == Post.modelName,
-                          let post = mutationEvent.element.model as? Post
-                    else
-                    { break }
-
-                    if post.title == newPost.title {
-                        initExpectation.fulfill()
-                    }
+            .filter { $0.eventName == HubPayload.EventName.DataStore.outboxMutationProcessed }
+            .map { DataStoreHubEvent(payload: $0) }
+            .compactMap(extractPost)
+            .sink { post in
+                if post.title == newPost.title {
+                    Task { await initExpectation.fulfill() }
+                }
 
-                    if post.title.hasPrefix(titlePrefix) {
-                        expectedResult.insert(post.title)
-                    }
+                if post.title.hasPrefix(titlePrefix) {
+                    expectedResult.insert(post.title)
+                }
 
-                    if expectedResult.count == parallelSize {
-                        parallelExpectation.fulfill()
-                    }
-                default: break
+                if expectedResult.count == parallelSize {
+                    Task { await parallelExpectation.fulfill() }
                 }
             }
         try await Amplify.DataStore.save(newPost)
-        wait(for: [initExpectation], timeout: 5)
+        await AsyncTesting.waitForExpectations([initExpectation], timeout: 10)
         try await Task.sleep(seconds: 1)
 
         for post in posts {
             Task {
                 try? await Amplify.DataStore.save(post)
             }
         }
-
-        wait(for: [parallelExpectation], timeout: Double(parallelSize))
+        await AsyncTesting.waitForExpectations([parallelExpectation], timeout: Double(parallelSize))
         cancellable.cancel()
         XCTAssertEqual(expectedResult, Set(posts.map(\.title)))
     }

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/TestSupport/HubEventsIntegrationTestBase.swift

@@ -33,11 +33,18 @@ class HubEventsIntegrationTestBase: XCTestCase {
     // swiftlint:enable force_try
     // swiftlint:enable force_cast
 
-    override func setUp() {
-        super.setUp()
+    override func setUp() async throws {
+        try await super.setUp()
         continueAfterFailure = false
     }
 
+    override func tearDown() async throws {
+        storageAdapter.delete(untypedModelType: ModelSyncMetadata.self, withIdentifier: ModelIdentifier<ModelSyncMetadata, ModelIdentifierFormat.Default>.makeDefault(id:"Post")) { _ in }
+        storageAdapter.delete(untypedModelType: ModelSyncMetadata.self, withIdentifier: ModelIdentifier<ModelSyncMetadata, ModelIdentifierFormat.Default>.makeDefault(id:"Comment")) { _ in }
+        await Amplify.reset()
+        try await Task.sleep(seconds: 1)
+    }
+
     func startAmplify(withModels models: AmplifyModelRegistration) async {
         do {
             let amplifyConfig = try TestConfigHelper.retrieveAmplifyConfiguration(forResource: Self.amplifyConfigurationFile)
@@ -51,11 +58,4 @@ class HubEventsIntegrationTestBase: XCTestCase {
             return
         }
     }
-
-    override func tearDown() async throws {
-        print("Amplify reset")
-        storageAdapter.delete(untypedModelType: ModelSyncMetadata.self, withIdentifier: ModelIdentifier<ModelSyncMetadata, ModelIdentifierFormat.Default>.makeDefault(id:"Post")) { _ in }
-        storageAdapter.delete(untypedModelType: ModelSyncMetadata.self, withIdentifier: ModelIdentifier<ModelSyncMetadata, ModelIdentifierFormat.Default>.makeDefault(id:"Comment")) { _ in }
-        await Amplify.reset()
-    }
 }

AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginIntegrationTests/TestSupport/SyncEngineIntegrationTestBase.swift

@@ -32,13 +32,23 @@ class SyncEngineIntegrationTestBase: DataStoreTestBase {
     // swiftlint:enable force_try
     // swiftlint:enable force_cast
 
+    override func setUp() async throws {
+        try await super.setUp()
+        continueAfterFailure = false
+    }
+
+    override func tearDown() async throws {
+        try await super.tearDown()
+        try await stopDataStore()
+        await Amplify.reset()
+        try await Task.sleep(seconds: 1)
+    }
+
     func setUp(
         withModels models: AmplifyModelRegistration,
         logLevel: LogLevel = .error,
         dataStoreConfiguration: DataStoreConfiguration? = nil
     ) async {
-        continueAfterFailure = false
-
         Amplify.Logging.logLevel = logLevel
 
         do {
@@ -55,11 +65,6 @@ class SyncEngineIntegrationTestBase: DataStoreTestBase {
         }
     }
 
-    override func tearDown() async throws {
-        try await Amplify.DataStore.stop()
-        await Amplify.reset()
-    }
-    
     func stopDataStore() async throws {
         try await Amplify.DataStore.stop()
     }
@@ -91,8 +96,9 @@ class SyncEngineIntegrationTestBase: DataStoreTestBase {
     }
 
     private func startAmplifyAndWait(for eventName: String) async throws {
-        let eventReceived = expectation(description: "DataStore \(eventName) event")
+        try startAmplify()
 
+        let eventReceived = expectation(description: "DataStore \(eventName) event")
         var token: UnsubscribeToken!
         token = Amplify.Hub.listen(to: .dataStore,
                                    eventName: eventName) { _ in
@@ -105,11 +111,10 @@ class SyncEngineIntegrationTestBase: DataStoreTestBase {
             return
         }
 
-        try startAmplify()
-        // try await Amplify.DataStore.start()
-        try await deleteMutationEvents()
+        try await Amplify.DataStore.start()
+        await waitForExpectations(timeout: 10.0)
 
-        await waitForExpectations(timeout: 100.0)
+        try await deleteMutationEvents()
     }
 
     func deleteMutationEvents() async throws {