Test useContext in pure, forwardRef, and PureComponent

packages/react-reconciler/src/__tests__/ReactNewContext-test.internal.js

@@ -57,6 +57,25 @@ describe('ReactNewContext', () => {
         return render(contextValue);
       },
   );
+  sharedContextTests('useContext inside forwardRef component', Context =>
+    React.forwardRef(function Consumer(props, ref) {
+      const observedBits = props.unstable_observedBits;
+      const contextValue = useContext(Context, observedBits);
+      const render = props.children;
+      return render(contextValue);
+    }),
+  );
+  sharedContextTests(
+    'useContext inside memoized function component',
+    Context =>
+      React.memo(function Consumer(props) {
+        const observedBits = props.unstable_observedBits;
+        const contextValue = useContext(Context, observedBits);
+        const render = props.children;
+        return render(contextValue);
+      }),
+    true,
+  );
   sharedContextTests(
     'readContext(Context) inside class component',
     Context =>
@@ -70,7 +89,7 @@ describe('ReactNewContext', () => {
       },
   );
 
-  function sharedContextTests(label, getConsumer) {
+  function sharedContextTests(label, getConsumer, isPure) {
     describe(`reading context with ${label}`, () => {
       it('simple mount and update', () => {
         const Context = React.createContext(1);
@@ -855,46 +874,48 @@ describe('ReactNewContext', () => {
         expect(ReactNoop.getChildren()).toEqual([span(2), span(2)]);
       });
 
-      // Context consumer bails out on propagating "deep" updates when `value` hasn't changed.
-      // However, it doesn't bail out from rendering if the component above it re-rendered anyway.
-      // If we bailed out on referential equality, it would be confusing that you
-      // can call this.setState(), but an autobound render callback "blocked" the update.
-      // https://github.com/facebook/react/pull/12470#issuecomment-376917711
-      it('consumer does not bail out if there were no bailouts above it', () => {
-        const Context = React.createContext(0);
-        const Consumer = getConsumer(Context);
-
-        class App extends React.Component {
-          state = {
-            text: 'hello',
-          };
-
-          renderConsumer = context => {
-            ReactNoop.yield('App#renderConsumer');
-            return <span prop={this.state.text} />;
-          };
-
-          render() {
-            ReactNoop.yield('App');
-            return (
-              <Context.Provider value={this.props.value}>
-                <Consumer>{this.renderConsumer}</Consumer>
-              </Context.Provider>
-            );
+      if (!isPure) {
+        // Context consumer bails out on propagating "deep" updates when `value` hasn't changed.
+        // However, it doesn't bail out from rendering if the component above it re-rendered anyway.
+        // If we bailed out on referential equality, it would be confusing that you
+        // can call this.setState(), but an autobound render callback "blocked" the update.
+        // https://github.com/facebook/react/pull/12470#issuecomment-376917711
+        it('consumer does not bail out if there were no bailouts above it', () => {
+          const Context = React.createContext(0);
+          const Consumer = getConsumer(Context);
+
+          class App extends React.Component {
+            state = {
+              text: 'hello',
+            };
+
+            renderConsumer = context => {
+              ReactNoop.yield('App#renderConsumer');
+              return <span prop={this.state.text} />;
+            };
+
+            render() {
+              ReactNoop.yield('App');
+              return (
+                <Context.Provider value={this.props.value}>
+                  <Consumer>{this.renderConsumer}</Consumer>
+                </Context.Provider>
+              );
+            }
           }
-        }
 
-        // Initial mount
-        let inst;
-        ReactNoop.render(<App value={1} ref={ref => (inst = ref)} />);
-        expect(ReactNoop.flush()).toEqual(['App', 'App#renderConsumer']);
-        expect(ReactNoop.getChildren()).toEqual([span('hello')]);
+          // Initial mount
+          let inst;
+          ReactNoop.render(<App value={1} ref={ref => (inst = ref)} />);
+          expect(ReactNoop.flush()).toEqual(['App', 'App#renderConsumer']);
+          expect(ReactNoop.getChildren()).toEqual([span('hello')]);
 
-        // Update
-        inst.setState({text: 'goodbye'});
-        expect(ReactNoop.flush()).toEqual(['App', 'App#renderConsumer']);
-        expect(ReactNoop.getChildren()).toEqual([span('goodbye')]);
-      });
+          // Update
+          inst.setState({text: 'goodbye'});
+          expect(ReactNoop.flush()).toEqual(['App', 'App#renderConsumer']);
+          expect(ReactNoop.getChildren()).toEqual([span('goodbye')]);
+        });
+      }
 
       // This is a regression case for https://github.com/facebook/react/issues/12389.
       it('does not run into an infinite loop', () => {