Improve tests that deal with microtasks

packages/internal-test-utils/ReactInternalTestUtils.js

@@ -218,6 +218,30 @@ ${diff(expectedLog, actualLog)}
   throw error;
 }
 
+export async function waitForDiscrete(expectedLog) {
+  assertYieldsWereCleared(SchedulerMock);
+
+  // Create the error object before doing any async work, to get a better
+  // stack trace.
+  const error = new Error();
+  Error.captureStackTrace(error, waitForDiscrete);
+
+  // Wait until end of current task/microtask.
+  await waitForMicrotasks();
+
+  const actualLog = SchedulerMock.unstable_clearLog();
+  if (equals(actualLog, expectedLog)) {
+    return;
+  }
+
+  error.message = `
+Expected sequence of events did not occur.
+
+${diff(expectedLog, actualLog)}
+`;
+  throw error;
+}
+
 export function assertLog(expectedLog) {
   const actualLog = SchedulerMock.unstable_clearLog();
   if (equals(actualLog, expectedLog)) {

packages/react-dom/src/__tests__/ReactDOMSafariMicrotaskBug-test.js

@@ -13,25 +13,33 @@ let React;
 
 let ReactDOMClient;
 let act;
+let assertLog;
+let Scheduler;
 
 describe('ReactDOMSafariMicrotaskBug-test', () => {
   let container;
-  let flushMicrotasksPrematurely;
+  let overrideQueueMicrotask;
+  let flushFakeMicrotasks;
 
   beforeEach(() => {
     // In Safari, microtasks don't always run on clean stack.
     // This setup crudely approximates it.
     // In reality, the sync flush happens when an iframe is added to the page.
     // https://github.com/facebook/react/issues/22459
-    let queue = [];
-    window.queueMicrotask = function (cb) {
-      queue.push(cb);
+    const originalQueueMicrotask = queueMicrotask;
+    overrideQueueMicrotask = false;
+    const fakeMicrotaskQueue = [];
+    global.queueMicrotask = cb => {
+      if (overrideQueueMicrotask) {
+        fakeMicrotaskQueue.push(cb);
+      } else {
+        originalQueueMicrotask(cb);
+      }
     };
-    flushMicrotasksPrematurely = function () {
-      while (queue.length > 0) {
-        const prevQueue = queue;
-        queue = [];
-        prevQueue.forEach(cb => cb());
+    flushFakeMicrotasks = () => {
+      while (fakeMicrotaskQueue.length > 0) {
+        const cb = fakeMicrotaskQueue.shift();
+        cb();
       }
     };
 
@@ -40,6 +48,8 @@ describe('ReactDOMSafariMicrotaskBug-test', () => {
     React = require('react');
     ReactDOMClient = require('react-dom/client');
     act = require('internal-test-utils').act;
+    assertLog = require('internal-test-utils').assertLog;
+    Scheduler = require('scheduler');
 
     document.body.appendChild(container);
   });
@@ -55,10 +65,14 @@ describe('ReactDOMSafariMicrotaskBug-test', () => {
       return (
         <div
           ref={() => {
+            overrideQueueMicrotask = true;
             if (!ran) {
               ran = true;
               setState(1);
-              flushMicrotasksPrematurely();
+              flushFakeMicrotasks();
+              Scheduler.log(
+                'Content at end of ref callback: ' + container.textContent,
+              );
             }
           }}>
           {state}
@@ -69,6 +83,7 @@ describe('ReactDOMSafariMicrotaskBug-test', () => {
     await act(() => {
       root.render(<Foo />);
     });
+    assertLog(['Content at end of ref callback: 0']);
     expect(container.textContent).toBe('1');
   });
 
@@ -78,8 +93,12 @@ describe('ReactDOMSafariMicrotaskBug-test', () => {
       return (
         <button
           onClick={() => {
+            overrideQueueMicrotask = true;
             setState(1);
-            flushMicrotasksPrematurely();
+            flushFakeMicrotasks();
+            Scheduler.log(
+              'Content at end of click handler: ' + container.textContent,
+            );
           }}>
           {state}
         </button>
@@ -95,6 +114,11 @@ describe('ReactDOMSafariMicrotaskBug-test', () => {
         new MouseEvent('click', {bubbles: true}),
       );
     });
+    // This causes the update to flush earlier than usual. This isn't the ideal
+    // behavior but we use this test to document it. The bug is Safari's, not
+    // ours, so we just do our best to not crash even though the behavior isn't
+    // completely correct.
+    assertLog(['Content at end of click handler: 1']);
     expect(container.textContent).toBe('1');
   });
 });

packages/react-dom/src/events/plugins/__tests__/ChangeEventPlugin-test.js

@@ -16,6 +16,7 @@ let ReactFeatureFlags;
 let Scheduler;
 let act;
 let waitForAll;
+let waitForDiscrete;
 let assertLog;
 
 const setUntrackedChecked = Object.getOwnPropertyDescriptor(
@@ -65,6 +66,7 @@ describe('ChangeEventPlugin', () => {
 
     const InternalTestUtils = require('internal-test-utils');
     waitForAll = InternalTestUtils.waitForAll;
+    waitForDiscrete = InternalTestUtils.waitForDiscrete;
     assertLog = InternalTestUtils.assertLog;
 
     container = document.createElement('div');
@@ -730,8 +732,7 @@ describe('ChangeEventPlugin', () => {
       );
 
       // Flush microtask queue.
-      await null;
-      assertLog(['render: ']);
+      await waitForDiscrete(['render: ']);
       expect(input.value).toBe('');
     });
 

packages/react-reconciler/src/__tests__/ReactFlushSyncNoAggregateError-test.js

@@ -5,6 +5,9 @@ let act;
 let assertLog;
 let waitForThrow;
 
+let overrideQueueMicrotask;
+let flushFakeMicrotasks;
+
 // TODO: Migrate tests to React DOM instead of React Noop
 
 describe('ReactFlushSync (AggregateError not available)', () => {
@@ -13,6 +16,26 @@ describe('ReactFlushSync (AggregateError not available)', () => {
 
     global.AggregateError = undefined;
 
+    // When AggregateError is not available, the errors are rethrown in a
+    // microtask. This is an implementation detail but we want to test it here
+    // so override the global one.
+    const originalQueueMicrotask = queueMicrotask;
+    overrideQueueMicrotask = false;
+    const fakeMicrotaskQueue = [];
+    global.queueMicrotask = cb => {
+      if (overrideQueueMicrotask) {
+        fakeMicrotaskQueue.push(cb);
+      } else {
+        originalQueueMicrotask(cb);
+      }
+    };
+    flushFakeMicrotasks = () => {
+      while (fakeMicrotaskQueue.length > 0) {
+        const cb = fakeMicrotaskQueue.shift();
+        cb();
+      }
+    };
+
     React = require('react');
     ReactNoop = require('react-noop-renderer');
     Scheduler = require('scheduler');
@@ -47,6 +70,8 @@ describe('ReactFlushSync (AggregateError not available)', () => {
     const aahh = new Error('AAHH!');
     const nooo = new Error('Noooooooooo!');
 
+    // Override the global queueMicrotask so we can test the behavior.
+    overrideQueueMicrotask = true;
     let error;
     try {
       ReactNoop.flushSync(() => {
@@ -70,10 +95,15 @@ describe('ReactFlushSync (AggregateError not available)', () => {
     // AggregateError is not available, React throws the first error, then
     // throws the remaining errors in separate tasks.
     expect(error).toBe(aahh);
+
     // TODO: Currently the remaining error is rethrown in an Immediate Scheduler
     // task, but this may change to a timer or microtask in the future. The
     // exact mechanism is an implementation detail; they just need to be logged
     // in the order the occurred.
+
+    // This will start throwing if we change it to rethrow in a microtask.
+    flushFakeMicrotasks();
+
     await waitForThrow(nooo);
   });
 });