From 2e0986fee87b2989437518df2d75bc468e202409 Mon Sep 17 00:00:00 2001
From: Joe Haddad <joe.haddad@zeit.co>
Date: Tue, 1 Sep 2020 13:43:44 -0400
Subject: [PATCH] [test] Update hydration marker for React 17 (#16756)

---
 packages/next/client/index.tsx | 18 ++++++++++--------
 test/acceptance/helpers.js     |  4 ----
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/packages/next/client/index.tsx b/packages/next/client/index.tsx
index 128e49063e9b5..3bf5a3b3f62c2 100644
--- a/packages/next/client/index.tsx
+++ b/packages/next/client/index.tsx
@@ -151,14 +151,6 @@ class Container extends React.Component<{
         }
       )
     }
-
-    if (process.env.__NEXT_TEST_MODE) {
-      window.__NEXT_HYDRATED = true
-
-      if (window.__NEXT_HYDRATED_CB) {
-        window.__NEXT_HYDRATED_CB()
-      }
-    }
   }
 
   componentDidUpdate() {
@@ -719,5 +711,15 @@ function Root({
   // We use `useLayoutEffect` to guarantee the callback is executed
   // as soon as React flushes the update.
   React.useLayoutEffect(() => callback(), [callback])
+  if (process.env.__NEXT_TEST_MODE) {
+    // eslint-disable-next-line react-hooks/rules-of-hooks
+    React.useEffect(() => {
+      window.__NEXT_HYDRATED = true
+
+      if (window.__NEXT_HYDRATED_CB) {
+        window.__NEXT_HYDRATED_CB()
+      }
+    }, [])
+  }
   return children as React.ReactElement
 }
diff --git a/test/acceptance/helpers.js b/test/acceptance/helpers.js
index 7379af1c2fae9..dd30e780f1b0b 100644
--- a/test/acceptance/helpers.js
+++ b/test/acceptance/helpers.js
@@ -37,10 +37,6 @@ export async function sandbox(
     env: { __NEXT_TEST_WITH_DEVTOOL: 1 },
   })
   const browser = await webdriver(appPort, '/')
-
-  // Slow down tests a bit to ensure application is ready:
-  await new Promise((resolve) => setTimeout(resolve, 750))
-
   return [
     {
       sandboxDirectory,