add setCollection support to Onyx.update()

lib/Onyx.ts

@@ -26,6 +26,7 @@ import type {
     OnyxUpdate,
     OnyxValue,
     OnyxInput,
+    OnyxMethodMap,
 } from './types';
 import OnyxUtils from './OnyxUtils';
 import logMessages from './logMessages';
@@ -599,7 +600,7 @@ function updateSnapshots(data: OnyxUpdate[]) {
 function update(data: OnyxUpdate[]): Promise<void> {
     // First, validate the Onyx object is in the format we expect
     data.forEach(({onyxMethod, key, value}) => {
-        if (![OnyxUtils.METHOD.CLEAR, OnyxUtils.METHOD.SET, OnyxUtils.METHOD.MERGE, OnyxUtils.METHOD.MERGE_COLLECTION, OnyxUtils.METHOD.MULTI_SET].includes(onyxMethod)) {
+        if (!Object.values(OnyxUtils.METHOD).includes(onyxMethod)) {
             throw new Error(`Invalid onyxMethod ${onyxMethod} in Onyx update.`);
         }
         if (onyxMethod === OnyxUtils.METHOD.MULTI_SET) {
@@ -636,18 +637,14 @@ function update(data: OnyxUpdate[]): Promise<void> {
     let clearPromise: Promise<void> = Promise.resolve();
 
     data.forEach(({onyxMethod, key, value}) => {
-        switch (onyxMethod) {
-            case OnyxUtils.METHOD.SET:
-                enqueueSetOperation(key, value);
-                break;
-            case OnyxUtils.METHOD.MERGE:
-                enqueueMergeOperation(key, value);
-                break;
-            case OnyxUtils.METHOD.MERGE_COLLECTION: {
+        const handlers: Record<OnyxMethodMap[keyof OnyxMethodMap], (k: typeof key, v: typeof value) => void> = {
+            [OnyxUtils.METHOD.SET]: enqueueSetOperation,
+            [OnyxUtils.METHOD.MERGE]: enqueueMergeOperation,
+            [OnyxUtils.METHOD.MERGE_COLLECTION]: () => {
                 const collection = value as Collection<CollectionKey, unknown, unknown>;
                 if (!OnyxUtils.isValidNonEmptyCollectionForMerge(collection)) {
                     Logger.logInfo('mergeCollection enqueued within update() with invalid or empty value. Skipping this operation.');
-                    break;
+                    return;
                 }
 
                 // Confirm all the collection keys belong to the same parent
@@ -656,18 +653,15 @@ function update(data: OnyxUpdate[]): Promise<void> {
                     const mergedCollection: OnyxInputKeyValueMapping = collection;
                     collectionKeys.forEach((collectionKey) => enqueueMergeOperation(collectionKey, mergedCollection[collectionKey]));
                 }
-
-                break;
-            }
-            case OnyxUtils.METHOD.MULTI_SET:
-                Object.entries(value).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue));
-                break;
-            case OnyxUtils.METHOD.CLEAR:
+            },
+            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollection(k, v as Collection<CollectionKey, unknown, unknown>)),
+            [OnyxUtils.METHOD.MULTI_SET]: (k, v) => Object.entries(v as Partial<OnyxInputKeyValueMapping>).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue)),
+            [OnyxUtils.METHOD.CLEAR]: () => {
                 clearPromise = clear();
-                break;
-            default:
-                break;
-        }
+            },
+        };
+
+        handlers[onyxMethod](key, value);
     });
 
     // Group all the collection-related keys and update each collection in a single `mergeCollection` call.

lib/OnyxUtils.ts

@@ -1418,4 +1418,6 @@ const OnyxUtils = {
     getEvictionBlocklist,
 };
 
+export type {OnyxMethod};
+
 export default OnyxUtils;

lib/types.ts

@@ -2,6 +2,7 @@ import type {Merge} from 'type-fest';
 import type {BuiltIns} from 'type-fest/source/internal';
 import type OnyxUtils from './OnyxUtils';
 import type {WithOnyxInstance, WithOnyxState} from './withOnyx/types';
+import type {OnyxMethod} from './OnyxUtils';
 
 /**
  * Utility type that excludes `null` from the type `TValue`.
@@ -395,41 +396,46 @@ type OnyxMergeInput<TKey extends OnyxKey> = OnyxInput<TKey>;
  */
 type OnyxMergeCollectionInput<TKey extends OnyxKey, TMap = object> = Collection<TKey, NonNullable<OnyxInput<TKey>>, TMap>;
 
+type OnyxMethodMap = typeof OnyxUtils.METHOD;
+
+// Maps onyx methods to their corresponding value types
+type OnyxMethodValueMap = {
+    [OnyxUtils.METHOD.SET]: {
+        key: OnyxKey;
+        value: OnyxSetInput<OnyxKey>;
+    };
+    [OnyxUtils.METHOD.MULTI_SET]: {
+        key: OnyxKey;
+        value: OnyxMultiSetInput;
+    };
+    [OnyxUtils.METHOD.MERGE]: {
+        key: OnyxKey;
+        value: OnyxMergeInput<OnyxKey>;
+    };
+    [OnyxUtils.METHOD.CLEAR]: {
+        key: OnyxKey;
+        value?: undefined;
+    };
+    [OnyxUtils.METHOD.MERGE_COLLECTION]: {
+        key: CollectionKeyBase;
+        value: OnyxMergeCollectionInput<CollectionKeyBase>;
+    };
+    [OnyxUtils.METHOD.SET_COLLECTION]: {
+        key: CollectionKeyBase;
+        value: OnyxMergeCollectionInput<CollectionKeyBase>;
+    };
+};
+
 /**
- * Represents different kinds of updates that can be passed to `Onyx.update()` method. It is a discriminated union of
- * different update methods (`SET`, `MERGE`, `MERGE_COLLECTION`), each with their own key and value structure.
+ * OnyxUpdate type includes all onyx methods used in OnyxMethodValueMap.
+ * If a new method is added to OnyxUtils.METHOD constant, it must be added to OnyxMethodValueMap type.
+ * Otherwise it will show static type errors.
  */
-type OnyxUpdate =
-    | {
-          [TKey in OnyxKey]:
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.SET;
-                    key: TKey;
-                    value: OnyxSetInput<TKey>;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.MULTI_SET;
-                    key: TKey;
-                    value: OnyxMultiSetInput;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.MERGE;
-                    key: TKey;
-                    value: OnyxMergeInput<TKey>;
-                }
-              | {
-                    onyxMethod: typeof OnyxUtils.METHOD.CLEAR;
-                    key: TKey;
-                    value?: undefined;
-                };
-      }[OnyxKey]
-    | {
-          [TKey in CollectionKeyBase]: {
-              onyxMethod: typeof OnyxUtils.METHOD.MERGE_COLLECTION;
-              key: TKey;
-              value: OnyxMergeCollectionInput<TKey>;
-          };
-      }[CollectionKeyBase];
+type OnyxUpdate = {
+    [Method in OnyxMethod]: {
+        onyxMethod: Method;
+    } & OnyxMethodValueMap[Method];
+}[OnyxMethod];
 
 /**
  * Represents the options used in `Onyx.init()` method.
@@ -507,6 +513,8 @@ export type {
     OnyxMultiSetInput,
     OnyxMergeInput,
     OnyxMergeCollectionInput,
+    OnyxMethod,
+    OnyxMethodMap,
     OnyxUpdate,
     OnyxValue,
     Selector,

tests/unit/onyxTest.ts

@@ -1340,7 +1340,7 @@ describe('Onyx', () => {
                 return waitForPromisesToResolve();
             })
             .then(() => {
-                expect(testKeyValue).toEqual(undefined);
+                expect(testKeyValue).toBeUndefined();
             });
     });
 
@@ -1807,5 +1807,103 @@ describe('Onyx', () => {
                 });
             });
         });
+
+        it('should properly handle setCollection operations in update()', () => {
+            const routeA = `${ONYX_KEYS.COLLECTION.ROUTES}A`;
+            const routeB = `${ONYX_KEYS.COLLECTION.ROUTES}B`;
+            const routeC = `${ONYX_KEYS.COLLECTION.ROUTES}C`;
+
+            let routesCollection: unknown;
+
+            connection = Onyx.connect({
+                key: ONYX_KEYS.COLLECTION.ROUTES,
+                initWithStoredValues: false,
+                callback: (value) => {
+                    routesCollection = value;
+                },
+                waitForCollectionCallback: true,
+            });
+
+            return Onyx.mergeCollection(ONYX_KEYS.COLLECTION.ROUTES, {
+                [routeA]: {name: 'Route A'},
+                [routeB]: {name: 'Route B'},
+                [routeC]: {name: 'Route C'},
+            } as GenericCollection)
+                .then(() => {
+                    return Onyx.update([
+                        {
+                            onyxMethod: Onyx.METHOD.SET_COLLECTION,
+                            key: ONYX_KEYS.COLLECTION.ROUTES,
+                            value: {
+                                [routeA]: {name: 'New Route A'},
+                                [routeB]: {name: 'New Route B'},
+                            },
+                        },
+                    ]);
+                })
+                .then(() => {
+                    expect(routesCollection).toEqual({
+                        [routeA]: {name: 'New Route A'},
+                        [routeB]: {name: 'New Route B'},
+                    });
+                });
+        });
+
+        it('should handle mixed operations with setCollection in update()', () => {
+            const routeA = `${ONYX_KEYS.COLLECTION.ROUTES}A`;
+            const routeB = `${ONYX_KEYS.COLLECTION.ROUTES}B`;
+            const testKey = ONYX_KEYS.TEST_KEY;
+            let routesCollection: unknown;
+
+            connection = Onyx.connect({
+                key: ONYX_KEYS.COLLECTION.ROUTES,
+                initWithStoredValues: false,
+                callback: (value) => {
+                    routesCollection = value;
+                },
+                waitForCollectionCallback: true,
+            });
+
+            let testKeyValue: unknown;
+            Onyx.connect({
+                key: testKey,
+                callback: (value) => {
+                    testKeyValue = value;
+                },
+            });
+
+            return Onyx.mergeCollection(ONYX_KEYS.COLLECTION.ROUTES, {
+                [routeA]: {name: 'Route A'},
+                [routeB]: {name: 'Route B'},
+            } as GenericCollection)
+                .then(() => {
+                    return Onyx.update([
+                        {
+                            onyxMethod: Onyx.METHOD.SET,
+                            key: testKey,
+                            value: 'test value',
+                        },
+                        {
+                            onyxMethod: Onyx.METHOD.SET_COLLECTION,
+                            key: ONYX_KEYS.COLLECTION.ROUTES,
+                            value: {
+                                [routeA]: {name: 'Final Route A'},
+                            },
+                        },
+                        {
+                            onyxMethod: Onyx.METHOD.MERGE,
+                            key: testKey,
+                            value: 'merged value',
+                        },
+                    ]);
+                })
+                .then(() => {
+                    expect(routesCollection).toEqual({
+                        [routeA]: {name: 'Final Route A'},
+                    });
+
+                    expect(testKeyValue).toBe('merged value');
+                });
+        });
     });
 });