From 57a008c800bdb0d6e55eb7288e5504745059d583 Mon Sep 17 00:00:00 2001
From: Douglas Armstrong <dougarmstrong@fb.com>
Date: Wed, 2 Jun 2021 15:21:08 -0700
Subject: [PATCH] Provide new value to onSet() handler when triggered by async
 default

Summary:
When an atom effect's `onSet()` handler is triggered due to having a default which is an async `Promise` that resolves the handler should be able to get the proper new value instead of a `DefaultValue` placeholder.

This addresses customer-reported issue #1050

Differential Revision: D28850148

fbshipit-source-id: 0f19d2179e7f3d522887550af32b87247f979f26
---
 src/recoil_values/Recoil_atom.js              | 16 +++----
 .../__tests__/Recoil_atom-test.js             | 45 ++++++++++++++++++-
 typescript/index.d.ts                         |  2 +-
 typescript/tests.ts                           |  4 +-
 4 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/src/recoil_values/Recoil_atom.js b/src/recoil_values/Recoil_atom.js
index d379aa325..3c48016bf 100644
--- a/src/recoil_values/Recoil_atom.js
+++ b/src/recoil_values/Recoil_atom.js
@@ -128,9 +128,7 @@ export type AtomEffect<T> = ({
 
   // Subscribe callbacks to events.
   // Atom effect observers are called before global transaction observers
-  onSet: (
-    (newValue: T | DefaultValue, oldValue: T | DefaultValue) => void,
-  ) => void,
+  onSet: ((newValue: T, oldValue: T | DefaultValue) => void) => void,
 }) => void | (() => void);
 
 export type AtomOptions<T> = $ReadOnly<{
@@ -293,9 +291,7 @@ function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
       };
       const resetSelf = effect => () => setSelf(effect)(DEFAULT_VALUE);
 
-      const onSet = effect => (
-        handler: (T | DefaultValue, T | DefaultValue) => void,
-      ) => {
+      const onSet = effect => (handler: (T, T | DefaultValue) => void) => {
         store.subscribeToTransactions(currentStore => {
           // eslint-disable-next-line prefer-const
           let {currentTree, previousTree} = currentStore.getState();
@@ -306,10 +302,10 @@ function baseAtom<T>(options: BaseAtomOptions<T>): RecoilState<T> {
             );
             previousTree = currentTree; // attempt to trundle on
           }
-          const newLoadable = currentTree.atomValues.get(key);
-          if (newLoadable == null || newLoadable.state === 'hasValue') {
-            const newValue: T | DefaultValue =
-              newLoadable != null ? newLoadable.contents : DEFAULT_VALUE;
+          const newLoadable =
+            currentTree.atomValues.get(key) ?? defaultLoadable;
+          if (newLoadable.state === 'hasValue') {
+            const newValue: T = newLoadable.contents;
             const oldLoadable =
               previousTree.atomValues.get(key) ?? defaultLoadable;
             const oldValue: T | DefaultValue =
diff --git a/src/recoil_values/__tests__/Recoil_atom-test.js b/src/recoil_values/__tests__/Recoil_atom-test.js
index 5cdb81ab5..f527c40da 100644
--- a/src/recoil_values/__tests__/Recoil_atom-test.js
+++ b/src/recoil_values/__tests__/Recoil_atom-test.js
@@ -260,7 +260,7 @@ describe('Effects', () => {
           setSelf('INIT');
           // This only fires on the reset action, not the default promise resolving
           onSet(newValue => {
-            expect(newValue).toBeInstanceOf(DefaultValue);
+            expect(newValue).toBe('RESOLVE');
           });
         },
       ],
@@ -443,6 +443,47 @@ describe('Effects', () => {
     expect(onSetForSameEffect).toHaveBeenCalledTimes(0);
   });
 
+  testRecoil('set default promise', async () => {
+    let setValue = 'RESOLVE_DEFAULT';
+    const onSetHandler = jest.fn(newValue => {
+      expect(newValue).toBe(setValue);
+    });
+
+    let resolveDefault;
+    const myAtom = atom({
+      key: 'atom effect default promise',
+      default: new Promise(resolve => {
+        resolveDefault = resolve;
+      }),
+      effects_UNSTABLE: [
+        ({onSet}) => {
+          onSet(onSetHandler);
+        },
+      ],
+    });
+
+    const [ReadsWritesAtom, set, reset] = componentThatReadsAndWritesAtom(
+      myAtom,
+    );
+    const c = renderElements(<ReadsWritesAtom />);
+    expect(c.textContent).toEqual('loading');
+
+    act(() => resolveDefault?.('RESOLVE_DEFAULT'));
+    await flushPromisesAndTimers();
+    expect(c.textContent).toEqual('"RESOLVE_DEFAULT"');
+    expect(onSetHandler).toHaveBeenCalledTimes(1);
+
+    setValue = 'SET';
+    act(() => set('SET'));
+    expect(c.textContent).toEqual('"SET"');
+    expect(onSetHandler).toHaveBeenCalledTimes(2);
+
+    setValue = 'RESOLVE_DEFAULT';
+    act(reset);
+    expect(c.textContent).toEqual('"RESOLVE_DEFAULT"');
+    expect(onSetHandler).toHaveBeenCalledTimes(3);
+  });
+
   testRecoil(
     'when setSelf is called in onSet, then onSet is not triggered again',
     () => {
@@ -556,7 +597,7 @@ describe('Effects', () => {
             }),
           );
           onSet(value => {
-            expect(value).toBeInstanceOf(DefaultValue);
+            expect(value).toBe('DEFAULT');
             validated = true;
           });
         },
diff --git a/typescript/index.d.ts b/typescript/index.d.ts
index 9b6e42a3c..d0fadb0c5 100644
--- a/typescript/index.d.ts
+++ b/typescript/index.d.ts
@@ -80,7 +80,7 @@ export type AtomEffect<T> = (param: {
   // Subscribe callbacks to events.
   // Atom effect observers are called before global transaction observers
   onSet: (
-    param: (newValue: T | DefaultValue, oldValue: T | DefaultValue) => void,
+    param: (newValue: T, oldValue: T | DefaultValue) => void,
   ) => void,
 }) => void | (() => void);
 
diff --git a/typescript/tests.ts b/typescript/tests.ts
index fb14b46f2..e1c011b61 100644
--- a/typescript/tests.ts
+++ b/typescript/tests.ts
@@ -470,7 +470,7 @@ isRecoilValue(mySelector1);
         setSelf('a'); // $ExpectError
 
         onSet(val => {
-          val; // $ExpectType number | DefaultValue
+          val; // $ExpectType number
         });
         onSet('a'); // $ExpectError
 
@@ -496,7 +496,7 @@ isRecoilValue(mySelector1);
         setSelf('a'); // $ExpectError
 
         onSet(val => {
-          val; // $ExpectType number | DefaultValue
+          val; // $ExpectType number
         });
         onSet('a'); // $ExpectError