fix: Onyx.clear handle collections and regular keys with an underscore

lib/Onyx.ts

@@ -554,22 +554,27 @@ function mergeCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TK
  * @param keysToPreserve is a list of ONYXKEYS that should not be cleared with the rest of the data
  */
 function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
+    const defaultKeyStates = OnyxUtils.getDefaultKeyStates();
+    const initialKeys = Object.keys(defaultKeyStates);
+
     return OnyxUtils.getAllKeys()
-        .then((keys) => {
+        .then((cachedKeys) => {
             cache.clearNullishStorageKeys();
 
             const keysToBeClearedFromStorage: OnyxKey[] = [];
             const keyValuesToResetAsCollection: Record<OnyxKey, OnyxCollection<KeyValueMapping[OnyxKey]>> = {};
             const keyValuesToResetIndividually: KeyValueMapping = {};
 
+            const allKeys = new Set([...cachedKeys, ...initialKeys]);
+
             // The only keys that should not be cleared are:
             // 1. Anything specifically passed in keysToPreserve (because some keys like language preferences, offline
             //      status, or activeClients need to remain in Onyx even when signed out)
             // 2. Any keys with a default state (because they need to remain in Onyx as their default, and setting them
             //      to null would cause unknown behavior)
-            keys.forEach((key) => {
+            //   2.1 However, if a default key was explicitly set to null, we need to reset it to the default value
+            allKeys.forEach((key) => {
                 const isKeyToPreserve = keysToPreserve.includes(key);
-                const defaultKeyStates = OnyxUtils.getDefaultKeyStates();
                 const isDefaultKey = key in defaultKeyStates;
 
                 // If the key is being removed or reset to default:
@@ -581,8 +586,9 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                     const newValue = defaultKeyStates[key] ?? null;
                     if (newValue !== oldValue) {
                         cache.set(key, newValue);
-                        const collectionKey = key.substring(0, key.indexOf('_') + 1);
-                        if (collectionKey) {
+
+                        const collectionKey = OnyxUtils.getCollectionKey(key);
+                        if (OnyxUtils.isCollectionKey(collectionKey)) {
                             if (!keyValuesToResetAsCollection[collectionKey]) {
                                 keyValuesToResetAsCollection[collectionKey] = {};
                             }
@@ -611,7 +617,6 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                 updatePromises.push(OnyxUtils.scheduleNotifyCollectionSubscribers(key, value));
             });
 
-            const defaultKeyStates = OnyxUtils.getDefaultKeyStates();
             const defaultKeyValuePairs = Object.entries(
                 Object.keys(defaultKeyStates)
                     .filter((key) => !keysToPreserve.includes(key))

lib/OnyxUtils.ts

@@ -403,14 +403,16 @@ function isCollectionMemberKey<TCollectionKey extends CollectionKeyBase>(collect
  * @param key - The collection member key to split.
  * @returns A tuple where the first element is the collection part and the second element is the ID part.
  */
-function splitCollectionMemberKey<TKey extends CollectionKey>(key: TKey): [TKey extends `${infer Prefix}_${string}` ? `${Prefix}_` : never, string] {
+function splitCollectionMemberKey<TKey extends CollectionKey, CollectionKeyType = TKey extends `${infer Prefix}_${string}` ? `${Prefix}_` : never>(key: TKey): [CollectionKeyType, string] {
     const underscoreIndex = key.lastIndexOf('_');
 
     if (underscoreIndex === -1) {
         throw new Error(`Invalid ${key} key provided, only collection keys are allowed.`);
     }
 
-    return [key.substring(0, underscoreIndex + 1) as TKey extends `${infer Prefix}_${string}` ? `${Prefix}_` : never, key.substring(underscoreIndex + 1)];
+    const collectionKey = key.substring(0, underscoreIndex + 1) as CollectionKeyType;
+    const memberKey = key.substring(underscoreIndex + 1);
+    return [collectionKey, memberKey];
 }
 
 /**

tests/unit/onyxTest.ts

@@ -9,6 +9,8 @@ import type GenericCollection from '../utils/GenericCollection';
 const ONYX_KEYS = {
     TEST_KEY: 'test',
     OTHER_TEST: 'otherTest',
+    // Special case: this key is not a collection key, but it has an underscore in its name
+    KEY_WITH_UNDERSCORE: 'nvp_test',
     COLLECTION: {
         TEST_KEY: 'test_',
         TEST_CONNECT_COLLECTION: 'testConnectCollection_',
@@ -25,6 +27,7 @@ Onyx.init({
     keys: ONYX_KEYS,
     initialKeyStates: {
         [ONYX_KEYS.OTHER_TEST]: 42,
+        [ONYX_KEYS.KEY_WITH_UNDERSCORE]: 'default',
     },
 });
 
@@ -51,15 +54,18 @@ describe('Onyx', () => {
                 expect(keys.has(ONYX_KEYS.OTHER_TEST)).toBe(true);
                 return Onyx.set(ONYX_KEYS.OTHER_TEST, null);
             })
+            // Checks if cache value is removed.
             .then(() => {
-                // Checks if cache value is removed.
-                expect(cache.getAllKeys().size).toBe(0);
-
-                // When cache keys length is 0, we fetch the keys from storage.
+                expect(cache.get(ONYX_KEYS.OTHER_TEST)).toBeUndefined();
                 return OnyxUtils.getAllKeys();
             })
             .then((keys) => {
                 expect(keys.has(ONYX_KEYS.OTHER_TEST)).toBe(false);
+            })
+            // Expect to reset to initial key value when calling Onyx.clear()
+            .then(() => Onyx.clear())
+            .then(() => {
+                expect(cache.get(ONYX_KEYS.OTHER_TEST)).toBe(42);
             }));
 
     it('should set a simple key', () => {
@@ -167,36 +173,64 @@ describe('Onyx', () => {
             },
         });
 
-        let otherTestValue: unknown;
-        const mockCallback = jest.fn((value) => {
-            otherTestValue = value;
-        });
+        const mockCallback = jest.fn();
         const otherTestConnectionID = Onyx.connect({
             key: ONYX_KEYS.OTHER_TEST,
             callback: mockCallback,
         });
 
         return waitForPromisesToResolve()
+            .then(() => {
+                expect(mockCallback).toHaveBeenCalledTimes(1);
+                expect(mockCallback).toHaveBeenCalledWith(42, ONYX_KEYS.OTHER_TEST);
+                mockCallback.mockClear();
+            })
             .then(() => Onyx.set(ONYX_KEYS.TEST_KEY, 'test'))
             .then(() => {
                 expect(testKeyValue).toBe('test');
-                mockCallback.mockReset();
-                return Onyx.clear().then(waitForPromisesToResolve);
+                return Onyx.clear();
             })
+            .then(() => waitForPromisesToResolve())
             .then(() => {
                 // Test key should be cleared
                 expect(testKeyValue).toBeUndefined();
 
-                // Expect that the connection to a key with a default value wasn't cleared
+                // Expect that the connection to a key with a default value that wasn't changed is not called on clear
                 expect(mockCallback).toHaveBeenCalledTimes(0);
 
-                // Other test key should be returned to its default state
-                expect(otherTestValue).toBe(42);
-
                 return Onyx.disconnect(otherTestConnectionID);
             });
     });
 
+    it('should notify key subscribers that use a underscore in their name', () => {
+        const mockCallback = jest.fn();
+        connectionID = Onyx.connect({
+            key: ONYX_KEYS.KEY_WITH_UNDERSCORE,
+            callback: mockCallback,
+        });
+
+        return waitForPromisesToResolve()
+            .then(() => mockCallback.mockReset())
+            .then(() => Onyx.set(ONYX_KEYS.KEY_WITH_UNDERSCORE, 'test'))
+            .then(() => {
+                expect(mockCallback).toHaveBeenCalledTimes(1);
+                expect(mockCallback).toHaveBeenLastCalledWith('test', ONYX_KEYS.KEY_WITH_UNDERSCORE);
+                mockCallback.mockReset();
+                return Onyx.clear();
+            })
+            .then(() => {
+                expect(mockCallback).toHaveBeenCalledTimes(1);
+                expect(mockCallback).toHaveBeenCalledWith('default', ONYX_KEYS.KEY_WITH_UNDERSCORE);
+            })
+            .then(() => Onyx.set(ONYX_KEYS.KEY_WITH_UNDERSCORE, 'default'))
+            .then(() => mockCallback.mockReset())
+            .then(() => Onyx.set(ONYX_KEYS.KEY_WITH_UNDERSCORE, 'test'))
+            .then(() => {
+                expect(mockCallback).toHaveBeenCalledTimes(1);
+                expect(mockCallback).toHaveBeenCalledWith('test', ONYX_KEYS.KEY_WITH_UNDERSCORE);
+            });
+    });
+
     it('should not notify subscribers after they have disconnected', () => {
         let testKeyValue: unknown;
         connectionID = Onyx.connect({