From 3653f10f836c3b6a12cc1b70f36fcf9eb32597e2 Mon Sep 17 00:00:00 2001
From: Dan <dan.abramov@gmail.com>
Date: Mon, 20 May 2019 10:09:40 +0100
Subject: [PATCH 1/3] Test that state is not leaked between components

---
 .../src/__tests__/ReactFresh-test.internal.js | 90 +++++++++++++++++++
 1 file changed, 90 insertions(+)

diff --git a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
index 4aa0d30bfad6..f5a14dc96fc7 100644
--- a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
+++ b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
@@ -963,6 +963,96 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('does not leak state between components', () => {
+    if (__DEV__) {
+      const AppV1 = render(
+        () => {
+          function Hello1() {
+            const [val, setVal] = React.useState(0);
+            return (
+              <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+                {val}
+              </p>
+            );
+          }
+          __register__(Hello1, 'Hello1');
+          function Hello2() {
+            const [val, setVal] = React.useState(0);
+            return (
+              <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+                {val}
+              </p>
+            );
+          }
+          __register__(Hello2, 'Hello2');
+          function App({cond}) {
+            return cond ? <Hello1 /> : <Hello2 />;
+          }
+          __register__(App, 'App');
+          return App;
+        },
+        {cond: false},
+      );
+
+      // Bump the state before patching.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('blue');
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(el.textContent).toBe('1');
+
+      // Switch the condition, flipping inner content.
+      // This should reset the state.
+      render(() => AppV1, {cond: true});
+      const el2 = container.firstChild;
+      expect(el2).not.toBe(el);
+      expect(el2.textContent).toBe('0');
+      expect(el2.style.color).toBe('blue');
+
+      // Bump it again.
+      act(() => {
+        el2.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(el2.textContent).toBe('1');
+
+      // Perform a hot update for Hello only.
+      patch(() => {
+        function Hello1() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello1, 'Hello1');
+        function Hello2() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello2, 'Hello2');
+      });
+
+      // Assert the state was preserved but color changed.
+      expect(container.firstChild).toBe(el2);
+      expect(el2.textContent).toBe('1');
+      expect(el2.style.color).toBe('red');
+
+      // Flip the condition again.
+      render(() => AppV1, {cond: false});
+      const el3 = container.firstChild;
+      expect(el3).not.toBe(el2);
+      expect(el3.textContent).toBe('0');
+      expect(el3.style.color).toBe('red');
+    }
+  });
+
   it('can force remount by changing signature', () => {
     if (__DEV__) {
       let HelloV1 = render(() => {

From ec99af4740bfd3eaa9a74018c57377e43654765f Mon Sep 17 00:00:00 2001
From: Dan <dan.abramov@gmail.com>
Date: Mon, 20 May 2019 11:08:37 +0100
Subject: [PATCH 2/3] Support lazy and add Suspense tests

---
 .../src/__tests__/ReactFresh-test.internal.js | 591 ++++++++++++++++++
 .../src/ReactFiberBeginWork.js                |  11 +-
 2 files changed, 599 insertions(+), 3 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
index f5a14dc96fc7..42d9f8f60b3b 100644
--- a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
+++ b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
@@ -895,6 +895,597 @@ describe('ReactFresh', () => {
     }
   });
 
+  it('can preserve state for lazy after resolution', async () => {
+    if (__DEV__) {
+      let promise;
+      let OuterV1 = render(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+
+        const Outer = React.lazy(() => {
+          promise = new Promise(resolve => {
+            setTimeout(() => resolve({default: Hello}), 100);
+          });
+          return promise;
+        });
+        __register__(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+      jest.runAllTimers();
+      await promise;
+      Scheduler.flushAll();
+      expect(container.textContent).toBe('0');
+
+      // Bump the state before patching.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('blue');
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(el.textContent).toBe('1');
+
+      // Perform a hot update.
+      patch(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+      });
+
+      // Assert the state was preserved but color changed.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('red');
+
+      // Bump the state again.
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('2');
+      expect(el.style.color).toBe('red');
+
+      // Perform top-down renders with a stale type.
+      render(() => OuterV1);
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('2');
+      expect(el.style.color).toBe('red');
+
+      // Finally, a render with incompatible type should reset it.
+      render(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+
+        // Note: no lazy wrapper this time.
+        return Hello;
+      });
+
+      expect(container.firstChild).not.toBe(el);
+      const newEl = container.firstChild;
+      expect(newEl.textContent).toBe('0');
+      expect(newEl.style.color).toBe('blue');
+    }
+  });
+
+  it('can patch lazy before resolution', async () => {
+    if (__DEV__) {
+      let promise;
+      render(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+
+        const Outer = React.lazy(() => {
+          promise = new Promise(resolve => {
+            setTimeout(() => resolve({default: Hello}), 100);
+          });
+          return promise;
+        });
+        __register__(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+
+      // Perform a hot update.
+      patch(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+      });
+
+      jest.runAllTimers();
+      await promise;
+      Scheduler.flushAll();
+
+      // Expect different color on initial mount.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('red');
+
+      // Bump state.
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('red');
+
+      // Test another reload.
+      patch(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'orange'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('orange');
+    }
+  });
+
+  it('can patch lazy(forwardRef) before resolution', async () => {
+    if (__DEV__) {
+      let promise;
+      render(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.forwardRef(renderHello);
+        __register__(Hello, 'Hello');
+
+        const Outer = React.lazy(() => {
+          promise = new Promise(resolve => {
+            setTimeout(() => resolve({default: Hello}), 100);
+          });
+          return promise;
+        });
+        __register__(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+
+      // Perform a hot update.
+      patch(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.forwardRef(renderHello);
+        __register__(Hello, 'Hello');
+      });
+
+      jest.runAllTimers();
+      await promise;
+      Scheduler.flushAll();
+
+      // Expect different color on initial mount.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('red');
+
+      // Bump state.
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('red');
+
+      // Test another reload.
+      patch(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'orange'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.forwardRef(renderHello);
+        __register__(Hello, 'Hello');
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('orange');
+    }
+  });
+
+  it('can patch lazy(memo) before resolution', async () => {
+    if (__DEV__) {
+      let promise;
+      render(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.memo(renderHello);
+        __register__(Hello, 'Hello');
+
+        const Outer = React.lazy(() => {
+          promise = new Promise(resolve => {
+            setTimeout(() => resolve({default: Hello}), 100);
+          });
+          return promise;
+        });
+        __register__(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+
+      // Perform a hot update.
+      patch(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.memo(renderHello);
+        __register__(Hello, 'Hello');
+      });
+
+      jest.runAllTimers();
+      await promise;
+      Scheduler.flushAll();
+
+      // Expect different color on initial mount.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('red');
+
+      // Bump state.
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('red');
+
+      // Test another reload.
+      patch(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'orange'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.memo(renderHello);
+        __register__(Hello, 'Hello');
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('orange');
+    }
+  });
+
+  it('can patch lazy(memo(forwardRef)) before resolution', async () => {
+    if (__DEV__) {
+      let promise;
+      render(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.memo(React.forwardRef(renderHello));
+        __register__(Hello, 'Hello');
+
+        const Outer = React.lazy(() => {
+          promise = new Promise(resolve => {
+            setTimeout(() => resolve({default: Hello}), 100);
+          });
+          return promise;
+        });
+        __register__(Outer, 'Outer');
+
+        function App() {
+          return (
+            <React.Suspense fallback={<p>Loading</p>}>
+              <Outer />
+            </React.Suspense>
+          );
+        }
+
+        return App;
+      });
+
+      expect(container.textContent).toBe('Loading');
+
+      // Perform a hot update.
+      patch(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.memo(React.forwardRef(renderHello));
+        __register__(Hello, 'Hello');
+      });
+
+      jest.runAllTimers();
+      await promise;
+      Scheduler.flushAll();
+
+      // Expect different color on initial mount.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('red');
+
+      // Bump state.
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('red');
+
+      // Test another reload.
+      patch(() => {
+        function renderHello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'orange'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        const Hello = React.memo(React.forwardRef(renderHello));
+        __register__(Hello, 'Hello');
+      });
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('orange');
+    }
+  });
+
+  it('can patch both trees while suspense is displaying the fallback', async () => {
+    if (__DEV__) {
+      const AppV1 = render(
+        () => {
+          function Hello({children}) {
+            const [val, setVal] = React.useState(0);
+            return (
+              <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+                {children} {val}
+              </p>
+            );
+          }
+          __register__(Hello, 'Hello');
+
+          function Never() {
+            throw new Promise(resolve => {});
+          }
+
+          function App({shouldSuspend}) {
+            return (
+              <React.Suspense fallback={<Hello>Fallback</Hello>}>
+                <Hello>Content</Hello>
+                {shouldSuspend && <Never />}
+              </React.Suspense>
+            );
+          }
+
+          return App;
+        },
+        {shouldSuspend: false},
+      );
+
+      // We start with just the primary tree.
+      expect(container.childNodes.length).toBe(1);
+      const primaryChild = container.firstChild;
+      expect(primaryChild.textContent).toBe('Content 0');
+      expect(primaryChild.style.color).toBe('blue');
+      expect(primaryChild.style.display).toBe('');
+
+      // Bump primary content state.
+      act(() => {
+        primaryChild.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.childNodes.length).toBe(1);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('blue');
+      expect(primaryChild.style.display).toBe('');
+
+      // Perform a hot update.
+      patch(() => {
+        function Hello({children}) {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'green'}} onClick={() => setVal(val + 1)}>
+              {children} {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+      });
+      expect(container.childNodes.length).toBe(1);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('green');
+      expect(primaryChild.style.display).toBe('');
+
+      // Now force the tree to suspend.
+      render(() => AppV1, {shouldSuspend: true});
+
+      // Expect to see two trees, one of them is hidden.
+      expect(container.childNodes.length).toBe(2);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      const fallbackChild = container.childNodes[1];
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('green');
+      expect(primaryChild.style.display).toBe('none');
+      expect(fallbackChild.textContent).toBe('Fallback 0');
+      expect(fallbackChild.style.color).toBe('green');
+      expect(fallbackChild.style.display).toBe('');
+
+      // Bump fallback state.
+      act(() => {
+        fallbackChild.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(container.childNodes.length).toBe(2);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      expect(container.childNodes[1]).toBe(fallbackChild);
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('green');
+      expect(primaryChild.style.display).toBe('none');
+      expect(fallbackChild.textContent).toBe('Fallback 1');
+      expect(fallbackChild.style.color).toBe('green');
+      expect(fallbackChild.style.display).toBe('');
+
+      // Perform a hot update.
+      patch(() => {
+        function Hello({children}) {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {children} {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+      });
+
+      // Colors inside both trees should change:
+      expect(container.childNodes.length).toBe(2);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      expect(container.childNodes[1]).toBe(fallbackChild);
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('red');
+      expect(primaryChild.style.display).toBe('none');
+      expect(fallbackChild.textContent).toBe('Fallback 1');
+      expect(fallbackChild.style.color).toBe('red');
+      expect(fallbackChild.style.display).toBe('');
+
+      // Only primary tree should exist now:
+      render(() => AppV1, {shouldSuspend: false});
+      expect(container.childNodes.length).toBe(1);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('red');
+      expect(primaryChild.style.display).toBe('');
+
+      // Perform a hot update.
+      patch(() => {
+        function Hello({children}) {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'orange'}} onClick={() => setVal(val + 1)}>
+              {children} {val}
+            </p>
+          );
+        }
+        __register__(Hello, 'Hello');
+      });
+      expect(container.childNodes.length).toBe(1);
+      expect(container.childNodes[0]).toBe(primaryChild);
+      expect(primaryChild.textContent).toBe('Content 1');
+      expect(primaryChild.style.color).toBe('orange');
+      expect(primaryChild.style.display).toBe('');
+    }
+  });
+
   it('does not re-render ancestor components unnecessarily during a hot update', () => {
     if (__DEV__) {
       let appRenders = 0;
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index 71b931f6bba6..761c128de63f 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -1030,9 +1030,6 @@ function mountLazyComponent(
   // Cancel and resume right after we know the tag.
   cancelWorkTimer(workInProgress);
   let Component = readLazyComponentType(elementType);
-  if (__DEV__) {
-    // TODO: resolve type for hot reloading.
-  }
   // Store the unwrapped component in the type.
   workInProgress.type = Component;
   const resolvedTag = (workInProgress.tag = resolveLazyComponentTag(Component));
@@ -1043,6 +1040,9 @@ function mountLazyComponent(
     case FunctionComponent: {
       if (__DEV__) {
         validateFunctionComponentInDev(workInProgress, Component);
+        workInProgress.type = Component = resolveFunctionForHotReloading(
+          Component,
+        );
       }
       child = updateFunctionComponent(
         null,
@@ -1064,6 +1064,11 @@ function mountLazyComponent(
       break;
     }
     case ForwardRef: {
+      if (__DEV__) {
+        workInProgress.type = Component = resolveFunctionForHotReloading(
+          Component,
+        );
+      }
       child = updateForwardRef(
         null,
         workInProgress,

From 2fb5d0ed002d39b0a6b36f9f67601220c814b1bc Mon Sep 17 00:00:00 2001
From: Dan Abramov <dan.abramov@gmail.com>
Date: Mon, 20 May 2019 13:35:48 +0100
Subject: [PATCH 3/3] Nits

---
 .../src/__tests__/ReactFresh-test.internal.js | 99 +++++++++----------
 1 file changed, 47 insertions(+), 52 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
index 42d9f8f60b3b..6cf4189854d7 100644
--- a/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
+++ b/packages/react-dom/src/__tests__/ReactFresh-test.internal.js
@@ -897,7 +897,6 @@ describe('ReactFresh', () => {
 
   it('can preserve state for lazy after resolution', async () => {
     if (__DEV__) {
-      let promise;
       let OuterV1 = render(() => {
         function Hello() {
           const [val, setVal] = React.useState(0);
@@ -909,12 +908,12 @@ describe('ReactFresh', () => {
         }
         __register__(Hello, 'Hello');
 
-        const Outer = React.lazy(() => {
-          promise = new Promise(resolve => {
-            setTimeout(() => resolve({default: Hello}), 100);
-          });
-          return promise;
-        });
+        const Outer = React.lazy(
+          () =>
+            new Promise(resolve => {
+              setTimeout(() => resolve({default: Hello}), 100);
+            }),
+        );
         __register__(Outer, 'Outer');
 
         function App() {
@@ -929,9 +928,9 @@ describe('ReactFresh', () => {
       });
 
       expect(container.textContent).toBe('Loading');
-      jest.runAllTimers();
-      await promise;
-      Scheduler.flushAll();
+      await act(async () => {
+        jest.runAllTimers();
+      });
       expect(container.textContent).toBe('0');
 
       // Bump the state before patching.
@@ -969,7 +968,7 @@ describe('ReactFresh', () => {
       expect(el.textContent).toBe('2');
       expect(el.style.color).toBe('red');
 
-      // Perform top-down renders with a stale type.
+      // Perform a top-down render with a stale type.
       render(() => OuterV1);
       expect(container.firstChild).toBe(el);
       expect(el.textContent).toBe('2');
@@ -1000,7 +999,6 @@ describe('ReactFresh', () => {
 
   it('can patch lazy before resolution', async () => {
     if (__DEV__) {
-      let promise;
       render(() => {
         function Hello() {
           const [val, setVal] = React.useState(0);
@@ -1012,12 +1010,12 @@ describe('ReactFresh', () => {
         }
         __register__(Hello, 'Hello');
 
-        const Outer = React.lazy(() => {
-          promise = new Promise(resolve => {
-            setTimeout(() => resolve({default: Hello}), 100);
-          });
-          return promise;
-        });
+        const Outer = React.lazy(
+          () =>
+            new Promise(resolve => {
+              setTimeout(() => resolve({default: Hello}), 100);
+            }),
+        );
         __register__(Outer, 'Outer');
 
         function App() {
@@ -1046,9 +1044,9 @@ describe('ReactFresh', () => {
         __register__(Hello, 'Hello');
       });
 
-      jest.runAllTimers();
-      await promise;
-      Scheduler.flushAll();
+      await act(async () => {
+        jest.runAllTimers();
+      });
 
       // Expect different color on initial mount.
       const el = container.firstChild;
@@ -1083,7 +1081,6 @@ describe('ReactFresh', () => {
 
   it('can patch lazy(forwardRef) before resolution', async () => {
     if (__DEV__) {
-      let promise;
       render(() => {
         function renderHello() {
           const [val, setVal] = React.useState(0);
@@ -1096,12 +1093,12 @@ describe('ReactFresh', () => {
         const Hello = React.forwardRef(renderHello);
         __register__(Hello, 'Hello');
 
-        const Outer = React.lazy(() => {
-          promise = new Promise(resolve => {
-            setTimeout(() => resolve({default: Hello}), 100);
-          });
-          return promise;
-        });
+        const Outer = React.lazy(
+          () =>
+            new Promise(resolve => {
+              setTimeout(() => resolve({default: Hello}), 100);
+            }),
+        );
         __register__(Outer, 'Outer');
 
         function App() {
@@ -1131,9 +1128,9 @@ describe('ReactFresh', () => {
         __register__(Hello, 'Hello');
       });
 
-      jest.runAllTimers();
-      await promise;
-      Scheduler.flushAll();
+      await act(async () => {
+        jest.runAllTimers();
+      });
 
       // Expect different color on initial mount.
       const el = container.firstChild;
@@ -1169,7 +1166,6 @@ describe('ReactFresh', () => {
 
   it('can patch lazy(memo) before resolution', async () => {
     if (__DEV__) {
-      let promise;
       render(() => {
         function renderHello() {
           const [val, setVal] = React.useState(0);
@@ -1182,12 +1178,12 @@ describe('ReactFresh', () => {
         const Hello = React.memo(renderHello);
         __register__(Hello, 'Hello');
 
-        const Outer = React.lazy(() => {
-          promise = new Promise(resolve => {
-            setTimeout(() => resolve({default: Hello}), 100);
-          });
-          return promise;
-        });
+        const Outer = React.lazy(
+          () =>
+            new Promise(resolve => {
+              setTimeout(() => resolve({default: Hello}), 100);
+            }),
+        );
         __register__(Outer, 'Outer');
 
         function App() {
@@ -1217,9 +1213,9 @@ describe('ReactFresh', () => {
         __register__(Hello, 'Hello');
       });
 
-      jest.runAllTimers();
-      await promise;
-      Scheduler.flushAll();
+      await act(async () => {
+        jest.runAllTimers();
+      });
 
       // Expect different color on initial mount.
       const el = container.firstChild;
@@ -1255,7 +1251,6 @@ describe('ReactFresh', () => {
 
   it('can patch lazy(memo(forwardRef)) before resolution', async () => {
     if (__DEV__) {
-      let promise;
       render(() => {
         function renderHello() {
           const [val, setVal] = React.useState(0);
@@ -1268,12 +1263,12 @@ describe('ReactFresh', () => {
         const Hello = React.memo(React.forwardRef(renderHello));
         __register__(Hello, 'Hello');
 
-        const Outer = React.lazy(() => {
-          promise = new Promise(resolve => {
-            setTimeout(() => resolve({default: Hello}), 100);
-          });
-          return promise;
-        });
+        const Outer = React.lazy(
+          () =>
+            new Promise(resolve => {
+              setTimeout(() => resolve({default: Hello}), 100);
+            }),
+        );
         __register__(Outer, 'Outer');
 
         function App() {
@@ -1303,9 +1298,9 @@ describe('ReactFresh', () => {
         __register__(Hello, 'Hello');
       });
 
-      jest.runAllTimers();
-      await promise;
-      Scheduler.flushAll();
+      await act(async () => {
+        jest.runAllTimers();
+      });
 
       // Expect different color on initial mount.
       const el = container.firstChild;
@@ -1608,7 +1603,7 @@ describe('ReactFresh', () => {
       });
       expect(el2.textContent).toBe('1');
 
-      // Perform a hot update for Hello only.
+      // Perform a hot update for both inner components.
       patch(() => {
         function Hello1() {
           const [val, setVal] = React.useState(0);