fix(datastore-v1): observeQuery snapshot on .modelSynced event (#2367)

manaswi223 · web-flow · commit 2c4f85aa8c18 · 2022-10-05T12:53:27.000-04:00
* fix(datastore): observeQuery snapshot on .modelSynced event

* chore(datastore): fix integ tests on DataStoreObserveQueryTests

* chore(datastore-v1): address PR comments
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/AWSDataStorePlugin.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/AWSDataStorePlugin.swift
@@ -236,10 +236,10 @@ final public class AWSDataStorePlugin: DataStoreCategoryPlugin {
             dataStorePublisher.send(input: mutationEvent)
         case .modelSyncedEvent(let modelSyncedEvent):
             log.verbose("Emitting DataStore event: modelSyncedEvent \(modelSyncedEvent)")
+            dispatchedModelSyncedEvents[modelSyncedEvent.modelName]?.set(true)
             let modelSyncedEventPayload = HubPayload(eventName: HubPayload.EventName.DataStore.modelSynced,
                                                      data: modelSyncedEvent)
             Amplify.Hub.dispatch(to: .dataStore, payload: modelSyncedEventPayload)
-            dispatchedModelSyncedEvents[modelSyncedEvent.modelName]?.set(true)
         case .syncQueriesReadyEvent:
             log.verbose("[Lifecycle event 4]: syncQueriesReady")
             let syncQueriesReadyEventPayload = HubPayload(eventName: HubPayload.EventName.DataStore.syncQueriesReady)
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Subscribe/DataStoreObserveQueryOperation.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Subscribe/DataStoreObserveQueryOperation.swift
@@ -131,6 +131,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
     var currentItems: SortedList<M>
     var batchItemsChangedSink: AnyCancellable?
     var itemsChangedSink: AnyCancellable?
+    var modelSyncedEventSink: AnyCancellable?
 
     /// Internal publisher for `ObserveQueryPublisher` to pass events back to subscribers
     let passthroughPublisher: PassthroughSubject<DataStoreQuerySnapshot<M>, DataStoreError>
@@ -182,6 +183,10 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
         if let batchItemsChangedSink = batchItemsChangedSink {
             batchItemsChangedSink.cancel()
         }
+
+        if let modelSyncedEventSink = modelSyncedEventSink {
+            modelSyncedEventSink.cancel()
+        }
         passthroughPublisher.send(completion: .finished)
         super.cancel()
         finish()
@@ -198,6 +203,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
             self.currentItems.reset()
             self.itemsChangedSink = nil
             self.batchItemsChangedSink = nil
+            self.modelSyncedEventSink = nil
         }
     }
 
@@ -250,6 +256,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
                 switch queryResult {
                 case .success(let queriedModels):
                     currentItems.set(sortedModels: queriedModels)
+                    subscribeToModelSyncedEvent()
                     sendSnapshot()
                 case .failure(let error):
                     self.passthroughPublisher.send(completion: .failure(error))
@@ -284,6 +291,18 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
                   receiveValue: onItemChangeAfterSync(mutationEvent:))
     }
 
+    func subscribeToModelSyncedEvent() {
+        modelSyncedEventSink = Amplify.Hub.publisher(for: .dataStore).sink { event in
+            if event.eventName == HubPayload.EventName.DataStore.modelSynced,
+               let modelSyncedEvent = event.data as? ModelSyncedEvent,
+               modelSyncedEvent.modelName == self.modelSchema.name {
+                self.serialQueue.async {
+                    self.sendSnapshot()
+                }
+            }
+        }
+    }
+
     func filterByModelName(mutationEvent: MutationEvent) -> Bool {
         // Filter in the model when it matches the model name for this operation
         mutationEvent.modelName == modelSchema.name
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginIntegrationTests/DataStoreObserveQueryTests.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginIntegrationTests/DataStoreObserveQueryTests.swift
@@ -66,6 +66,50 @@ class DataStoreObserveQueryTests: SyncEngineIntegrationTestBase {
         sink.cancel()
     }
 
+    /// ObserveQuery API will eventually return a snapshot when sync state is toggled  from false to true.
+    ///  A `.modelSynced` event from the hub is internally received
+    ///
+    /// - Given: DataStore is cleared
+    /// - When:
+    ///    - ObserveQuery API is called to start the sync engine
+    ///    - A model is saved but not yet synced
+    /// - Then:
+    ///    - A query snapshot is received on `.modelSynced`
+    ///
+    func testObserveQueryWhenModelSyncedEvent() throws {
+        setUp(withModels: TestModelRegistration())
+        try startAmplify()
+        clearDataStore()
+        var snapshots = [DataStoreQuerySnapshot<Post>]()
+        var isObserveQueryReadyForTest = false
+        let observeQueryReadyForTest = expectation(description: "received query snapshot with .isSynced true")
+        let snapshotWithPost = expectation(description: "received first snapshot")
+        let post = Post(title: "title", content: "content", createdAt: .now())
+        let sink = Amplify.DataStore.observeQuery(for: Post.self).sink { completed in
+            switch completed {
+            case .finished:
+                break
+            case .failure(let error):
+                XCTFail("\(error)")
+            }
+        } receiveValue: { querySnapshot in
+            snapshots.append(querySnapshot)
+            if !isObserveQueryReadyForTest && querySnapshot.isSynced {
+                isObserveQueryReadyForTest = true
+                observeQueryReadyForTest.fulfill()
+            }
+            if querySnapshot.items.contains(where: { $0.id == post.id }) {
+                snapshotWithPost.fulfill()
+            }
+        }
+
+        wait(for: [observeQueryReadyForTest], timeout: 100)
+
+        _ = Amplify.DataStore.save(post)
+        wait(for: [snapshotWithPost], timeout: 100)
+        sink.cancel()
+    }
+
     /// Apply a query predicate "title begins with 'xyz'"
     ///
     /// - Given: DataStore is set up with an empty local store
@@ -83,6 +127,7 @@ class DataStoreObserveQueryTests: SyncEngineIntegrationTestBase {
         clearDataStore()
         var snapshots = [DataStoreQuerySnapshot<Post>]()
         let snapshotWithIsSynced = expectation(description: "query snapshot with isSynced true")
+        let receivedPostFromObserveQuery = expectation(description: "received Post")
         snapshotWithIsSynced.assertForOverFulfill = false
         var snapshotWithIsSyncedFulfilled = false
         let predicate = Post.keys.title.beginsWith("xyz")
@@ -94,18 +139,21 @@ class DataStoreObserveQueryTests: SyncEngineIntegrationTestBase {
                 XCTFail("\(error)")
             }
         } receiveValue: { querySnapshot in
-            if !snapshotWithIsSyncedFulfilled {
-                snapshots.append(querySnapshot)
-
-                if querySnapshot.isSynced {
-                    snapshotWithIsSyncedFulfilled = true
-                    snapshotWithIsSynced.fulfill()
+            snapshots.append(querySnapshot)
+            if !snapshotWithIsSyncedFulfilled && querySnapshot.isSynced {
+                snapshotWithIsSyncedFulfilled = true
+                snapshotWithIsSynced.fulfill()
+            } else if snapshotWithIsSyncedFulfilled {
+                if querySnapshot.items.count >= 4 && querySnapshot.items.contains(where: { post in
+                    post.title.contains("xyz")
+                }) {
+                    receivedPostFromObserveQuery.fulfill()
                 }
             }
         }
 
-        savePostAndWaitForSync(Post(title: "xyz", content: "content", createdAt: .now()))
-        wait(for: [snapshotWithIsSynced], timeout: 100)
+        savePostAndWaitForSync(Post(title: "xyz 4", content: "content", createdAt: .now()))
+        wait(for: [snapshotWithIsSynced, receivedPostFromObserveQuery], timeout: 200)
         XCTAssertTrue(snapshots.count >= 2)
         XCTAssertFalse(snapshots[0].isSynced)
         log.info("\(snapshots)")
@@ -213,7 +261,7 @@ class DataStoreObserveQueryTests: SyncEngineIntegrationTestBase {
     ///     - Delete a model that does NOT match the predicate. No snapshot is emitted
     func testPredicateWithCreateUpdateDelete() throws {
         setUp(withModels: TestModelRegistration(), logLevel: .info)
-        try startAmplify()
+        try startAmplifyAndWaitForReady()
 
         let testId = UUID().uuidString
         let postMatchPredicate = Post(title: "xyz 1", content: testId, createdAt: .now())
@@ -307,7 +355,7 @@ class DataStoreObserveQueryTests: SyncEngineIntegrationTestBase {
     ///
     func testSortWithCreateUpdateDelete() throws {
         setUp(withModels: TestModelRegistration(), logLevel: .info)
-        try startAmplify()
+        try startAmplifyAndWaitForReady()
 
         let testId = UUID().uuidString
         var snapshotCount = 0
diff --git a/AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Subscribe/DataStoreObserveQueryOperationTests.swift b/AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Subscribe/DataStoreObserveQueryOperationTests.swift
@@ -77,6 +77,128 @@ class DataStoreObserveQueryOperationTests: XCTestCase {
         sink.cancel()
     }
 
+    ///  ObserveQuery will send a single snapshot when the sync state toggles
+    ///  from false to true. The operation internally listens to `.modelSynced` event from
+    ///  the Hub.
+    ///
+    /// - Given: ObserveQuery has started and the first snapshot has been received.
+    /// - When:
+    ///    - modelSyncedEvent is sent to the Hub
+    /// - Then:
+    ///    - ObserveQuery will send a second snapshot
+    ///
+    func testGenerateSnapshotOnObserveQueryWhenModelSynced() throws {
+        let firstSnapshot = expectation(description: "first query snapshots")
+        let secondSnapshot = expectation(description: "second query snapshots")
+        let thirdSnapshot = expectation(description: "third query snapshot")
+        thirdSnapshot.isInverted = true
+        
+        var querySnapshots = [DataStoreQuerySnapshot<Post>]()
+        let dispatchedModelSyncedEvent = AtomicValue(initialValue: false)
+        let operation = AWSDataStoreObserveQueryOperation(
+            modelType: Post.self,
+            modelSchema: Post.schema,
+            predicate: nil,
+            sortInput: nil,
+            storageEngine: storageEngine,
+            dataStorePublisher: dataStorePublisher,
+            dataStoreConfiguration: .default,
+            dispatchedModelSyncedEvent: dispatchedModelSyncedEvent)
+
+        let sink = operation.publisher.sink { completed in
+            switch completed {
+            case .finished:
+                break
+            case .failure(let error):
+                XCTFail("Failed with error \(error)")
+            }
+        } receiveValue: { querySnapshot in
+            querySnapshots.append(querySnapshot)
+            if querySnapshots.count == 1 {
+                firstSnapshot.fulfill()
+            } else if querySnapshots.count == 2 {
+                secondSnapshot.fulfill()
+            } else if querySnapshots.count == 3 {
+                XCTFail("Should not receive third snapshot for a Model change")
+                thirdSnapshot.fulfill()
+            }
+        }
+        let queue = OperationQueue()
+        queue.addOperation(operation)
+        wait(for: [firstSnapshot], timeout: 5)
+
+        dispatchedModelSyncedEvent.set(true)
+        var modelSyncedEventBuilder = ModelSyncedEvent.Builder()
+        modelSyncedEventBuilder.modelName = Post.modelName
+        let modelSyncedEvent = modelSyncedEventBuilder.build()
+        let modelSyncedEventPayload = HubPayload(eventName: HubPayload.EventName.DataStore.modelSynced,
+                                                 data: modelSyncedEvent)
+        Amplify.Hub.dispatch(to: .dataStore, payload: modelSyncedEventPayload)
+        wait(for: [secondSnapshot], timeout: 10)
+
+        modelSyncedEventBuilder.modelName = "modelNameNotMatch"
+        let modelSyncedEventNotMatch = modelSyncedEventBuilder.build()
+        let modelSyncedEventNotMatchPayload = HubPayload(eventName: HubPayload.EventName.DataStore.modelSynced,
+                                                         data: modelSyncedEventNotMatch)
+        Amplify.Hub.dispatch(to: .dataStore, payload: modelSyncedEventNotMatchPayload)
+        wait(for: [thirdSnapshot], timeout: 10)
+
+        XCTAssertEqual(querySnapshots[0].items.count, 0)
+        XCTAssertEqual(querySnapshots[0].isSynced, false)
+        XCTAssertEqual(querySnapshots[1].isSynced, true)
+        sink.cancel()
+    }
+
+    /// ObserveQuery will send the first snapshot with 2 items when storage engine
+    /// is mocked to return 2 items.
+    ///
+    /// - Given: ObserveQuery starts
+    /// - When:
+    ///    -  ObserveQuery performs the initial query, two posts are queried through StorageEngine
+    /// - Then:
+    ///    - The items queried will return two posts in the first snapshot
+    ///
+    func testFirstSnapshotFromStorageQueryReturnsTwoPosts() throws {
+        let firstSnapshot = expectation(description: "firstSnapshot received")
+
+        var snapshots = [DataStoreQuerySnapshot<Post>]()
+        let dispatchedModelSyncedEvent = AtomicValue(initialValue: false)
+        let operation = AWSDataStoreObserveQueryOperation(
+            modelType: Post.self,
+            modelSchema: Post.schema,
+            predicate: nil,
+            sortInput: nil,
+            storageEngine: storageEngine,
+            dataStorePublisher: dataStorePublisher,
+            dataStoreConfiguration: .default,
+            dispatchedModelSyncedEvent: dispatchedModelSyncedEvent)
+        let post = Post(title: "model1",
+                        content: "content1",
+                        createdAt: .now())
+        storageEngine.responders[.query] = QueryResponder<Post>(callback: { _ in
+            return .success([post, post])
+        })
+
+        let sink = operation.publisher.sink { completed in
+            switch completed {
+            case .finished:
+                break
+            case .failure(let error):
+                XCTFail("Failed with error \(error)")
+            }
+        } receiveValue: { querySnapshot in
+            snapshots.append(querySnapshot)
+            if snapshots.count == 1 {
+                XCTAssertEqual(querySnapshot.items.count, 2)
+                firstSnapshot.fulfill()
+            }
+        }
+        let queue = OperationQueue()
+        queue.addOperation(operation)
+        wait(for: [firstSnapshot], timeout: 10)
+        sink.cancel()
+    }
+
     /// Multiple item changed observed will be returned in a single snapshot
     ///
     /// - Given:  The operation has started and the first query has completed.
