fix(DataStore): ObserveQuery InitialQuery empty modelIds set (#1663)

Author: lawmicha

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Subscribe/DataStoreObserveQueryOperation.swift

@@ -195,7 +195,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
                 self.observeQueryStarted = false
             }
             self.log.verbose("Resetting state")
-            self.currentItems.sortedModels.removeAll()
+            self.currentItems.reset()
             self.itemsChangedSink = nil
             self.batchItemsChangedSink = nil
         }
@@ -249,7 +249,7 @@ public class AWSDataStoreObserveQueryOperation<M: Model>: AsynchronousOperation,
 
                 switch queryResult {
                 case .success(let queriedModels):
-                    currentItems.sortedModels = queriedModels
+                    currentItems.set(sortedModels: queriedModels)
                     sendSnapshot()
                 case .failure(let error):
                     self.passthroughPublisher.send(completion: .failure(error))

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Subscribe/Support/SortedList.swift

@@ -8,8 +8,8 @@
 import Amplify
 
 class SortedList<ModelType: Model> {
-    var sortedModels: [ModelType]
-    var modelIds: Set<Model.Identifier>
+    private(set) var sortedModels: [ModelType]
+    private(set) var modelIds: Set<Model.Identifier>
     private let sortInput: [QuerySortDescriptor]?
     private let modelSchema: ModelSchema
 
@@ -20,10 +20,19 @@ class SortedList<ModelType: Model> {
         self.modelSchema = modelSchema
     }
 
+    /// Sets the internal data structures to the list of models. `sortedModels` should already be sorted to efficiently
+    /// update the internal instructures of `SortedList`. If `sortedModels` is not sorted, then the subsequent
+    /// methods used on the data structure will have adverse consequences, for example adding a model assumes the list
+    /// is already sorted (since it performs a binary search insertion).
+    func set(sortedModels: [ModelType]) {
+        self.sortedModels = sortedModels
+        modelIds = Set(sortedModels.map { $0.id })
+    }
+
     /// Apply the incoming `model` to the sorted array based on the mutation type. This logic accounts for duplicate
     /// events since identical events may have different sources (local and remote). When the mutation type is delete,
-    /// remove it if it exists in the array. When create/update and it has an sort order, then remove and add it back in the
-    /// correct sort order. if there is no sort order, replace it.
+    /// remove it if it exists in the array. When create/update and it has an sort order, then remove and add it back in
+    /// the correct sort order. if there is no sort order, replace it.
     /// Return `true` if something occured (added, replaced, deleted), otherwise `false`
     @discardableResult func apply(model: ModelType, mutationType: MutationEvent.MutationType) -> Bool {
         if mutationType == MutationEvent.MutationType.delete {
@@ -87,6 +96,12 @@ class SortedList<ModelType: Model> {
             modelIds.insert(model.id)
         }
     }
+
+    /// Removes the models in the list
+    func reset() {
+        sortedModels.removeAll()
+        modelIds.removeAll()
+    }
 }
 
 extension Array where Element: Model {

AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Subscribe/Support/SortedListTests.swift

@@ -16,6 +16,24 @@ import Combine
 
 class SortedListTests: XCTestCase {
 
+    func testSortedListSetAndReset() {
+        let posts = [createPost(id: "1", rating: 5.0),
+                     createPost(id: "2", rating: 5.0),
+                     createPost(id: "3", rating: 5.0),
+                     createPost(id: "4", rating: 5.0)]
+
+        let list = SortedList<Post>(sortInput: [QuerySortBy.ascending(Post.keys.rating).sortDescriptor],
+                                    modelSchema: Post.schema)
+        list.set(sortedModels: posts)
+
+        assertPosts(list.sortedModels, expectedIds: ["1", "2", "3", "4"])
+        XCTAssertEqual(list.modelIds.count, 4)
+
+        list.reset()
+        XCTAssertTrue(list.sortedModels.isEmpty)
+        XCTAssertTrue(list.modelIds.isEmpty)
+    }
+
     func testSortedListSingleSort() {
         let posts = [createPost(id: "1"),
                      createPost(id: "2"),
@@ -24,15 +42,18 @@ class SortedListTests: XCTestCase {
         let sortInput = [QuerySortBy.ascending(Post.keys.id).sortDescriptor]
         let list = SortedList<Post>(sortInput: sortInput,
                                     modelSchema: Post.schema)
-        list.sortedModels = posts
+
+        list.set(sortedModels: posts)
 
         // insert into the middle of the sorted list
         list.add(model: createPost(id: "3"), sortInputs: sortInput)
         assertPosts(list.sortedModels, expectedIds: ["1", "2", "3", "5", "6"])
+        XCTAssertEqual(list.modelIds.count, 5)
 
         // insert into the middle
         list.add(model: createPost(id: "4"), sortInputs: sortInput)
         assertPosts(list.sortedModels, expectedIds: ["1", "2", "3", "4", "5", "6"])
+        XCTAssertEqual(list.modelIds.count, 6)
 
         // insert at the beginning
         list.add(model: createPost(id: "0"), sortInputs: sortInput)
@@ -68,7 +89,7 @@ class SortedListTests: XCTestCase {
                          QuerySortBy.ascending(Post.keys.id).sortDescriptor]
         let list = SortedList<Post>(sortInput: sortInput,
                                     modelSchema: Post.schema)
-        list.sortedModels = posts
+        list.set(sortedModels: posts)
 
         // After id: "1", rating: 5.0
         list.add(model: createPost(id: "1", rating: 10.0), sortInputs: sortInput)
@@ -107,7 +128,7 @@ class SortedListTests: XCTestCase {
         let sortInput = [QuerySortBy.ascending(Post.keys.rating).sortDescriptor]
         let list = SortedList<Post>(sortInput: [QuerySortBy.ascending(Post.keys.rating).sortDescriptor],
                                     modelSchema: Post.schema)
-        list.sortedModels = posts
+        list.set(sortedModels: posts)
 
         // Since this is a binary search, the first index where the predicate returns `nil` is the middle index
         list.add(model: createPost(id: "5", rating: 5.0), sortInputs: sortInput)