From cf2b9e80a0b099c2f861660fd9cf589848106fa1 Mon Sep 17 00:00:00 2001
From: Jonathan Kim <jkimbo@gmail.com>
Date: Wed, 2 May 2018 17:10:25 +0100
Subject: [PATCH] Support forwardRef and Context in pretty-format (#6093)

---
 CHANGELOG.md                                  |  2 ++
 .../pretty-format/src/__tests__/react.test.js | 35 +++++++++++++++++++
 .../src/plugins/react_element.js              | 34 ++++++++++++------
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 729b99022783..335f3926ea70 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -125,6 +125,8 @@
   ([#5720](https://github.com/facebook/jest/pull/5720))
 * `[pretty-format]` Handle React fragments better
   ([#5816](https://github.com/facebook/jest/pull/5816))
+* `[pretty-format]` Handle formatting of `React.forwardRef` and `Context`
+  components ([#6093](https://github.com/facebook/jest/pull/6093))
 * `[jest-cli]` Switch collectCoverageFrom back to a string
   ([#5914](https://github.com/facebook/jest/pull/5914))
 * `[jest-regex-util]` Fix handling regex symbols in tests path on Windows
diff --git a/packages/pretty-format/src/__tests__/react.test.js b/packages/pretty-format/src/__tests__/react.test.js
index f136884305d9..f25d7b30187d 100644
--- a/packages/pretty-format/src/__tests__/react.test.js
+++ b/packages/pretty-format/src/__tests__/react.test.js
@@ -717,3 +717,38 @@ test('ReactTestComponent plugin highlights syntax with color from theme option',
     }),
   ).toMatchSnapshot();
 });
+
+test('supports forwardRef with a child', () => {
+  function Cat(props) {
+    return React.createElement('div', props, props.children);
+  }
+
+  expect(
+    // $FlowFixMe - https://github.com/facebook/flow/issues/6103
+    formatElement(React.createElement(React.forwardRef(Cat), null, 'mouse')),
+  ).toEqual('<ForwardRef(Cat)>\n  mouse\n</ForwardRef(Cat)>');
+});
+
+test('supports context Provider with a child', () => {
+  const {Provider} = React.createContext('test');
+
+  expect(
+    formatElement(
+      React.createElement(Provider, {value: 'test-value'}, 'child'),
+    ),
+  ).toEqual(
+    '<Context.Provider\n  value="test-value"\n>\n  child\n</Context.Provider>',
+  );
+});
+
+test('supports context Consumer with a child', () => {
+  const {Consumer} = React.createContext('test');
+
+  expect(
+    formatElement(
+      React.createElement(Consumer, null, () =>
+        React.createElement('div', null, 'child'),
+      ),
+    ),
+  ).toEqual('<Context.Consumer>\n  [Function anonymous]\n</Context.Consumer>');
+});
diff --git a/packages/pretty-format/src/plugins/react_element.js b/packages/pretty-format/src/plugins/react_element.js
index 1fd2be25b83d..78b73be8ae2e 100644
--- a/packages/pretty-format/src/plugins/react_element.js
+++ b/packages/pretty-format/src/plugins/react_element.js
@@ -19,6 +19,8 @@ import {
 const elementSymbol = Symbol.for('react.element');
 const fragmentSymbol = Symbol.for('react.fragment');
 const forwardRefSymbol = Symbol.for('react.forward_ref');
+const providerSymbol = Symbol.for('react.provider');
+const contextSymbol = Symbol.for('react.context');
 
 // Given element.props.children, or subtree during recursive traversal,
 // return flattened array of children.
@@ -34,22 +36,32 @@ const getChildren = (arg, children = []) => {
 };
 
 const getType = element => {
-  if (typeof element.type === 'string') {
-    return element.type;
+  const type = element.type;
+  if (typeof type === 'string') {
+    return type;
   }
-  if (typeof element.type === 'function') {
-    return element.type.displayName || element.type.name || 'Unknown';
+  if (typeof type === 'function') {
+    return type.displayName || type.name || 'Unknown';
   }
-  if (element.type === fragmentSymbol) {
+  if (type === fragmentSymbol) {
     return 'React.Fragment';
   }
-  if (element.type === forwardRefSymbol) {
-    const functionName =
-      element.type.render.displayName || element.type.render.name || '';
+  if (typeof type === 'object' && type !== null) {
+    if (type.$$typeof === providerSymbol) {
+      return 'Context.Provider';
+    }
 
-    return functionName !== ''
-      ? 'ForwardRef(' + functionName + ')'
-      : 'ForwardRef';
+    if (type.$$typeof === contextSymbol) {
+      return 'Context.Consumer';
+    }
+
+    if (type.$$typeof === forwardRefSymbol) {
+      const functionName = type.render.displayName || type.render.name || '';
+
+      return functionName !== ''
+        ? 'ForwardRef(' + functionName + ')'
+        : 'ForwardRef';
+    }
   }
   return 'UNDEFINED';
 };