Testing nits

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

@@ -289,16 +289,16 @@ describe('ReactDOMNativeEventHeuristic-test', () => {
     const target = React.createRef(null);
     function Foo() {
       const [count, setCount] = React.useState(0);
-
-      React.useEffect(() => {
-        Scheduler.unstable_yieldValue(count);
-      }, [count]);
+      const countRef = React.useRef(-1);
 
       React.useLayoutEffect(() => {
+        countRef.current = count;
         target.current.onclick = () => {
-          setCount(c => c + 1);
-          // Now update again, this should be batched.
-          setCount(c => c + 1);
+          setCount(countRef.current + 1);
+          // Now update again. If these updates are batched, then this should be
+          // a no-op, because we didn't re-render yet and `countRef` hasn't
+          // been mutated.
+          setCount(countRef.current + 1);
         };
       });
       return <div ref={target}>Count: {count}</div>;
@@ -307,7 +307,6 @@ describe('ReactDOMNativeEventHeuristic-test', () => {
     await act(async () => {
       root.render(<Foo />);
     });
-    expect(Scheduler).toHaveYielded([0]);
     expect(container.textContent).toEqual('Count: 0');
 
     // Ignore act warning. We can't use act because it forces batched updates.
@@ -316,16 +315,16 @@ describe('ReactDOMNativeEventHeuristic-test', () => {
     const pressEvent = document.createEvent('Event');
     pressEvent.initEvent('click', true, true);
     dispatchAndSetCurrentEvent(target.current, pressEvent);
-
-    expect(container.textContent).toEqual('Count: 2');
-    if (gate(flags => flags.enableDiscreteEventFlushingChange)) {
-      // When enableDiscreteEventFlushingChange is enabled, we don't flush
-      // discrete callbacks synchronously at the end of the event.
-      // TODO: I believe this flag is no longer relevant, because we flush the
-      // updates in a microtask, anyway.
-      expect(Scheduler).toFlushAndYield([2]);
-    } else {
-      expect(Scheduler).toHaveYielded([2]);
-    }
+    // If this is 2, that means the `setCount` calls were not batched.
+    expect(container.textContent).toEqual('Count: 1');
+
+    // Assert that the `act` warnings were the only ones that fired.
+    expect(console.error).toHaveBeenCalledTimes(2);
+    expect(console.error.calls.argsFor(0)[0]).toContain(
+      'was not wrapped in act',
+    );
+    expect(console.error.calls.argsFor(1)[0]).toContain(
+      'was not wrapped in act',
+    );
   });
 });