From 209e445e0536532539a202646ee043a025dd4237 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A1bio=20Henriques?= <contact@fabiohenriques.dev>
Date: Mon, 14 Oct 2024 22:28:59 +0100
Subject: [PATCH] Update useOnyx result if there is a pending Onyx.clear() task

---
 lib/Onyx.ts               |  4 +++-
 lib/useOnyx.ts            |  4 ++--
 tests/unit/useOnyxTest.ts | 17 +++++++++++------
 3 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/lib/Onyx.ts b/lib/Onyx.ts
index 27738b7c..6275e78b 100644
--- a/lib/Onyx.ts
+++ b/lib/Onyx.ts
@@ -449,7 +449,7 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
     const defaultKeyStates = OnyxUtils.getDefaultKeyStates();
     const initialKeys = Object.keys(defaultKeyStates);
 
-    return OnyxUtils.getAllKeys()
+    const promise = OnyxUtils.getAllKeys()
         .then((cachedKeys) => {
             cache.clearNullishStorageKeys();
 
@@ -537,6 +537,8 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                 });
         })
         .then(() => undefined);
+
+    return cache.captureTask('clear', promise) as Promise<void>;
 }
 
 function updateSnapshots(data: OnyxUpdate[]) {
diff --git a/lib/useOnyx.ts b/lib/useOnyx.ts
index dd8c1c0b..f2857eac 100644
--- a/lib/useOnyx.ts
+++ b/lib/useOnyx.ts
@@ -238,8 +238,8 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(key: TKey
 
         // If the previously cached value is different from the new value, we update both cached value
         // and the result to be returned by the hook.
-        // If the cache was set for the first time, we also update the cached value and the result.
-        const isCacheSetFirstTime = previousValueRef.current === null && hasCacheForKey;
+        // If the cache was set for the first time or we have a pending Onyx.clear() task, we also update the cached value and the result.
+        const isCacheSetFirstTime = previousValueRef.current === null && (hasCacheForKey || OnyxCache.hasPendingTask('clear'));
         if (isCacheSetFirstTime || !areValuesEqual) {
             previousValueRef.current = newValueRef.current;
 
diff --git a/tests/unit/useOnyxTest.ts b/tests/unit/useOnyxTest.ts
index e41499a6..6e74ee4a 100644
--- a/tests/unit/useOnyxTest.ts
+++ b/tests/unit/useOnyxTest.ts
@@ -150,15 +150,21 @@ describe('useOnyx', () => {
 
             Onyx.clear();
 
-            const {result} = renderHook(() => useOnyx(ONYXKEYS.TEST_KEY));
+            const {result: result1} = renderHook(() => useOnyx(ONYXKEYS.TEST_KEY));
+            const {result: result2} = renderHook(() => useOnyx(ONYXKEYS.TEST_KEY));
 
-            expect(result.current[0]).toBeUndefined();
-            expect(result.current[1].status).toEqual('loading');
+            expect(result1.current[0]).toBeUndefined();
+            expect(result1.current[1].status).toEqual('loaded');
+            expect(result2.current[0]).toBeUndefined();
+            expect(result2.current[1].status).toEqual('loaded');
 
+            Onyx.merge(ONYXKEYS.TEST_KEY, 'test2');
             await act(async () => waitForPromisesToResolve());
 
-            expect(result.current[0]).toBeUndefined();
-            expect(result.current[1].status).toEqual('loaded');
+            expect(result1.current[0]).toEqual('test2');
+            expect(result1.current[1].status).toEqual('loaded');
+            expect(result2.current[0]).toEqual('test2');
+            expect(result2.current[1].status).toEqual('loaded');
         });
 
         it('should return updated state when connecting to the same key after an Onyx.clear() call', async () => {
@@ -174,7 +180,6 @@ describe('useOnyx', () => {
             await act(async () => Onyx.clear());
 
             const {result: result2} = renderHook(() => useOnyx(ONYXKEYS.TEST_KEY));
-
             const {result: result3} = renderHook(() => useOnyx(ONYXKEYS.TEST_KEY));
 
             await act(async () => waitForPromisesToResolve());