expose a flag on the host config, move stuff around

packages/react-dom/src/client/ReactDOMHostConfig.js

@@ -39,6 +39,7 @@ import {
   DOCUMENT_FRAGMENT_NODE,
 } from '../shared/HTMLNodeType';
 import dangerousStyleValue from '../shared/dangerousStyleValue';
+import {canUseDOM} from 'shared/ExecutionEnvironment';
 
 import type {DOMContainer} from './ReactDOM';
 
@@ -647,3 +648,20 @@ export function didNotFindHydratableTextInstance(
     warnForInsertedHydratedText(parentInstance, text);
   }
 }
+
+// in a dom-like test environment, we want to warn if dispatchAction()
+// is called outside of a batchedUpdates/TestUtils.act(...) call.
+let ensureBatchedDispatchedActions = false;
+
+if (__DEV__) {
+  if (
+    canUseDOM &&
+    (navigator.userAgent.indexOf('Node.js') >= 0 ||
+      navigator.userAgent.indexOf('jsdom') >= 0)
+    // we should probably add puppeteer-like environments here too
+  ) {
+    ensureBatchedDispatchedActions = true;
+  }
+}
+
+export const shouldBatchDispatchedActions = ensureBatchedDispatchedActions;

packages/react-reconciler/src/ReactFiberHooks.js

@@ -31,17 +31,17 @@ import {
 } from './ReactHookEffectTags';
 import {
   scheduleWork,
-  getBatchingStatus,
+  ensureBatchingAndScheduleWork,
   computeExpirationForFiber,
   flushPassiveEffects,
   requestCurrentTime,
 } from './ReactFiberScheduler';
 
+import {shouldBatchDispatchedActions} from './ReactFiberHostConfig';
 import invariant from 'shared/invariant';
 import warning from 'shared/warning';
 import getComponentName from 'shared/getComponentName';
 import is from 'shared/objectIs';
-import {canUseDOM} from 'shared/ExecutionEnvironment';
 import {markWorkInProgressReceivedUpdate} from './ReactFiberBeginWork';
 
 type Update<S, A> = {
@@ -1129,66 +1129,13 @@ function dispatchAction<S, A>(
       }
     }
     if (__DEV__) {
-      ensureBatchingAndScheduleWork(fiber, expirationTime);
-    } else {
-      scheduleWork(fiber, expirationTime);
-    }
-  }
-}
-
-// in a dom-like test environment, we want to warn if dispatchAction()
-// is called outside of a batchedUpdates/TestUtils.act(...) call.
-
-let isNotDOMRenderer = undefined;
-// some folks initialize a jsdom, but still use TestRenderer.
-// So when dispatchAction is called first, we'll test the
-// fiber tree and set this value.
-
-let ensureBatchingAndScheduleWork = (
-  fiber: Fiber,
-  expirationTime: ExpirationTime,
-) => {
-  scheduleWork(fiber, expirationTime);
-};
-// we'll swap out the function definition when we learn more about the environment
-
-function ensureBatchingAndScheduleWorkImpl(
-  fiber: Fiber,
-  expirationTime: ExpirationTime,
-) {
-  if (__DEV__) {
-    if (isNotDOMRenderer === undefined) {
-      let f = fiber;
-      while (f.return) {
-        f = f.return;
-      }
-      if (f.stateNode.containerInfo instanceof HTMLElement) {
-        isNotDOMRenderer = false;
+      if (shouldBatchDispatchedActions === true) {
+        ensureBatchingAndScheduleWork(fiber, expirationTime);
       } else {
-        isNotDOMRenderer = true;
-        ensureBatchingAndScheduleWork = scheduleWork;
+        scheduleWork(fiber, expirationTime);
       }
+    } else {
+      scheduleWork(fiber, expirationTime);
     }
-    if (isNotDOMRenderer === false && getBatchingStatus() === false) {
-      warning(
-        false,
-        'It looks like you are in a test environment, trying to ' +
-          'set state outside of TestUtils.act(...). ' +
-          'This could lead to unexpected ui while testing. Use ' +
-          'ReactTestUtils.act(...) to batch your updates and remove this warning.',
-      );
-    }
-  }
-  scheduleWork(fiber, expirationTime);
-}
-
-if (__DEV__) {
-  if (
-    canUseDOM &&
-    (navigator.userAgent.indexOf('Node.js') >= 0 ||
-      navigator.userAgent.indexOf('jsdom') >= 0)
-    // we should probably add puppeteer-like environments here too
-  ) {
-    ensureBatchingAndScheduleWork = ensureBatchingAndScheduleWorkImpl;
   }
 }

packages/react-reconciler/src/ReactFiberScheduler.js

@@ -1789,8 +1789,22 @@ function scheduleWorkToRoot(fiber: Fiber, expirationTime): FiberRoot | null {
   return root;
 }
 
-export function getBatchingStatus(): boolean {
-  return isBatchingUpdates;
+export function ensureBatchingAndScheduleWork(
+  fiber: Fiber,
+  expirationTime: ExpirationTime,
+) {
+  if (__DEV__) {
+    if (isBatchingUpdates === false) {
+      warningWithoutStack(
+        false,
+        'It looks like you are in a test environment, trying to ' +
+          'set state outside of TestUtils.act(...). ' +
+          'This could lead to unexpected ui while testing. Use ' +
+          'ReactTestUtils.act(...) to batch your updates and remove this warning.',
+      );
+    }
+  }
+  scheduleWork(fiber, expirationTime);
 }
 
 function scheduleWork(fiber: Fiber, expirationTime: ExpirationTime) {