Add waitForCollectionCallback to keysChanged

lib/Onyx.js

@@ -260,6 +260,12 @@ function keysChanged(collectionKey, collection) {
         if (isSubscribedToCollectionKey) {
             if (_.isFunction(subscriber.callback)) {
                 const cachedCollection = getCachedCollection(collectionKey);
+
+                if (subscriber.waitForCollectionCallback) {
+                    subscriber.callback(cachedCollection);
+                    return;
+                }
+
                 _.each(collection, (data, dataKey) => {
                     subscriber.callback(cachedCollection[dataKey], dataKey);
                 });

tests/unit/onyxTest.js

@@ -473,8 +473,8 @@ describe('Onyx', () => {
         const valuesReceived = {};
         const mockCallback = jest.fn(data => valuesReceived[data.ID] = data.value);
 
-        // GIVEN some collection data
-        const collectionData = {
+        // GIVEN some initial collection data
+        const initialCollectionData = {
             test_connect_collection_1: {
                 ID: 123,
                 value: 'one',
@@ -488,23 +488,50 @@ describe('Onyx', () => {
                 value: 'three',
             },
         };
-        Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_CONNECT_COLLECTION, collectionData);
+
+        Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_CONNECT_COLLECTION, initialCollectionData);
         return waitForPromisesToResolve()
             .then(() => {
-                // WHEN we call Onyx.connect with the waitForCollectionCallback = true param
-                connectionID = Onyx.connect({
+                // WHEN we connect to that collection with waitForCollectionCallback = true
+                Onyx.connect({
                     key: ONYX_KEYS.COLLECTION.TEST_CONNECT_COLLECTION,
                     waitForCollectionCallback: true,
                     callback: mockCallback,
                 });
                 return waitForPromisesToResolve();
             })
             .then(() => {
-                // THEN the callback should be triggered only once
+                // THEN we expect the callback to be called only once and the initial stored value to be initialCollectionData
                 expect(mockCallback.mock.calls.length).toBe(1);
+                expect(mockCallback.mock.calls[0][0]).toEqual(initialCollectionData);
+            });
+    });
+
+    it('should return all collection keys as a single object when updating a collection key with waitForCollectionCallback = true', () => {
+        const valuesReceived = {};
+        const mockCallback = jest.fn(data => valuesReceived[data.ID] = data.value);
+        const collectionUpdate = {
+            test_connect_collection_3: {ID: 234, value: 'three'},
+            test_connect_collection_4: {ID: 123, value: 'four'},
+        };
 
-                // AND all the collection data should be returned as a single object
-                expect(mockCallback.mock.calls[0][0]).toEqual(collectionData);
+        // GIVEN an Onyx.connect call with waitForCollectionCallback=true
+        Onyx.connect({
+            key: ONYX_KEYS.COLLECTION.TEST_CONNECT_COLLECTION,
+            waitForCollectionCallback: true,
+            callback: mockCallback,
+        });
+        return waitForPromisesToResolve()
+            .then(() => {
+                // WHEN we update the collection, e.g. the API returns a response
+                Onyx.mergeCollection(ONYX_KEYS.COLLECTION.TEST_CONNECT_COLLECTION, collectionUpdate);
+                return waitForPromisesToResolve();
+            })
+            .then(() => {
+                // THEN we expect the callback to have called twice (once for the initial connect call and once for the collection update)
+                // AND the value should be equal to collectionUpdate
+                expect(mockCallback.mock.calls.length).toBe(2);
+                expect(mockCallback.mock.calls[1][0]).toEqual(collectionUpdate);
             });
     });
 });