Hardened logic around when and how to patch console methods

src/__tests__/console-test.js

@@ -4,10 +4,9 @@ describe('console', () => {
   let React;
   let ReactDOM;
   let act;
-  let enableConsole;
-  let disableConsole;
   let fakeConsole;
   let mockError;
+  let mockInfo;
   let mockLog;
   let mockWarn;
   let patchConsole;
@@ -15,8 +14,6 @@ describe('console', () => {
 
   beforeEach(() => {
     const Console = require('../backend/console');
-    enableConsole = Console.enable;
-    disableConsole = Console.disable;
     patchConsole = Console.patch;
     unpatchConsole = Console.unpatch;
 
@@ -37,25 +34,36 @@ describe('console', () => {
     // Patching the real console is too complicated,
     // because Jest itself has hooks into it as does our test env setup.
     mockError = jest.fn();
+    mockInfo = jest.fn();
     mockLog = jest.fn();
     mockWarn = jest.fn();
     fakeConsole = {
       error: mockError,
+      info: mockInfo,
       log: mockLog,
       warn: mockWarn,
     };
 
-    patchConsole(fakeConsole);
+    Console.dangerous_setTargetConsoleForTesting(fakeConsole);
+
+    patchConsole();
   });
 
   function normalizeCodeLocInfo(str) {
     return str && str.replace(/\(at .+?:\d+\)/g, '(at **)');
   }
 
+  it('should not patch console methods that do not receive component stacks', () => {
+    expect(fakeConsole.error).not.toBe(mockError);
+    expect(fakeConsole.info).toBe(mockInfo);
+    expect(fakeConsole.log).toBe(mockLog);
+    expect(fakeConsole.warn).not.toBe(mockWarn);
+  });
+
   it('should only patch the console once', () => {
     const { error, warn } = fakeConsole;
 
-    patchConsole(fakeConsole);
+    patchConsole();
 
     expect(fakeConsole.error).toBe(error);
     expect(fakeConsole.warn).toBe(warn);
@@ -89,30 +97,6 @@ describe('console', () => {
     expect(mockError.mock.calls[0][0]).toBe('error');
   });
 
-  it('should suppress console logging when disabled', () => {
-    disableConsole();
-    fakeConsole.log('log');
-    fakeConsole.warn('warn');
-    fakeConsole.error('error');
-    expect(mockLog).toHaveBeenCalledTimes(0);
-    expect(mockWarn).toHaveBeenCalledTimes(0);
-    expect(mockError).toHaveBeenCalledTimes(0);
-
-    enableConsole();
-    fakeConsole.log('log');
-    fakeConsole.warn('warn');
-    fakeConsole.error('error');
-    expect(mockLog).toHaveBeenCalledTimes(1);
-    expect(mockLog.mock.calls[0]).toHaveLength(1);
-    expect(mockLog.mock.calls[0][0]).toBe('log');
-    expect(mockWarn).toHaveBeenCalledTimes(1);
-    expect(mockWarn.mock.calls[0]).toHaveLength(1);
-    expect(mockWarn.mock.calls[0][0]).toBe('warn');
-    expect(mockError).toHaveBeenCalledTimes(1);
-    expect(mockError.mock.calls[0]).toHaveLength(1);
-    expect(mockError.mock.calls[0][0]).toBe('error');
-  });
-
   it('should not append multiple stacks', () => {
     const Child = () => {
       fakeConsole.warn('warn\n    in Child (at fake.js:123)');
@@ -330,7 +314,7 @@ describe('console', () => {
     expect(mockError.mock.calls[0]).toHaveLength(1);
     expect(mockError.mock.calls[0][0]).toBe('error');
 
-    patchConsole(fakeConsole);
+    patchConsole();
     act(() => ReactDOM.render(<Child />, document.createElement('div')));
 
     expect(mockWarn).toHaveBeenCalledTimes(2);

src/__tests__/inspectedElementContext-test.js

@@ -255,6 +255,74 @@ describe('InspectedElementContext', () => {
     done();
   });
 
+  it('should temporarily disable console logging when re-running a component to inspect its hooks', async done => {
+    let targetRenderCount = 0;
+
+    const errorSpy = ((console: any).error = jest.fn());
+    const infoSpy = ((console: any).info = jest.fn());
+    const logSpy = ((console: any).log = jest.fn());
+    const warnSpy = ((console: any).warn = jest.fn());
+
+    const Target = React.memo(props => {
+      targetRenderCount++;
+      console.error('error');
+      console.info('info');
+      console.log('log');
+      console.warn('warn');
+      React.useState(0);
+      return null;
+    });
+
+    const container = document.createElement('div');
+    await utils.actAsync(() =>
+      ReactDOM.render(<Target a={1} b="abc" />, container)
+    );
+
+    expect(targetRenderCount).toBe(1);
+    expect(errorSpy).toHaveBeenCalledTimes(1);
+    expect(errorSpy).toHaveBeenCalledWith('error');
+    expect(infoSpy).toHaveBeenCalledTimes(1);
+    expect(infoSpy).toHaveBeenCalledWith('info');
+    expect(logSpy).toHaveBeenCalledTimes(1);
+    expect(logSpy).toHaveBeenCalledWith('log');
+    expect(warnSpy).toHaveBeenCalledTimes(1);
+    expect(warnSpy).toHaveBeenCalledWith('warn');
+
+    const id = ((store.getElementIDAtIndex(0): any): number);
+
+    let inspectedElement = null;
+
+    function Suspender({ target }) {
+      const { getInspectedElement } = React.useContext(InspectedElementContext);
+      inspectedElement = getInspectedElement(target);
+      return null;
+    }
+
+    await utils.actAsync(
+      () =>
+        TestRenderer.create(
+          <Contexts
+            defaultSelectedElementID={id}
+            defaultSelectedElementIndex={1}
+          >
+            <React.Suspense fallback={null}>
+              <Suspender target={id} />
+            </React.Suspense>
+          </Contexts>
+        ),
+      false
+    );
+
+    expect(inspectedElement).not.toBe(null);
+    expect(targetRenderCount).toBe(2);
+    expect(errorSpy).toHaveBeenCalledTimes(1);
+    expect(infoSpy).toHaveBeenCalledTimes(1);
+    expect(logSpy).toHaveBeenCalledTimes(1);
+    expect(warnSpy).toHaveBeenCalledTimes(1);
+
+    done();
+  });
+
   it('should support simple data types', async done => {
     const Example = () => null;
 

src/backend/console.js

@@ -5,6 +5,8 @@ import describeComponentFrame from './describeComponentFrame';
 
 import type { Fiber, ReactRenderer } from './types';
 
+const APPEND_STACK_TO_METHODS = ['error', 'trace', 'warn'];
+
 const FRAME_REGEX = /\n {4}in /;
 
 const injectedRenderers: Map<
@@ -15,17 +17,29 @@ const injectedRenderers: Map<
   |}
 > = new Map();
 
-let isDisabled: boolean = false;
+let targetConsole: Object = console;
+let targetConsoleMethods = {};
+for (let method in console) {
+  targetConsoleMethods[method] = console[method];
+}
+
 let unpatchFn: null | (() => void) = null;
 
-export function disable(): void {
-  isDisabled = true;
-}
+// Enables e.g. Jest tests to inject a mock console object.
+export function dangerous_setTargetConsoleForTesting(
+  targetConsoleForTesting: Object
+): void {
+  targetConsole = targetConsoleForTesting;
 
-export function enable(): void {
-  isDisabled = false;
+  targetConsoleMethods = {};
+  for (let method in targetConsole) {
+    targetConsoleMethods[method] = console[method];
+  }
 }
 
+// v16 renderers should use this method to inject internals necessary to generate a component stack.
+// These internals will be used if the console is patched.
+// Injecting them separately allows the console to easily be patched or unpacted later (at runtime).
 export function registerRenderer(renderer: ReactRenderer): void {
   const { getCurrentFiber, findFiberByHostInstance, version } = renderer;
 
@@ -44,16 +58,18 @@ export function registerRenderer(renderer: ReactRenderer): void {
   }
 }
 
-export function patch(targetConsole?: Object = console): void {
+// Patches whitelisted console methods to append component stack for the current fiber.
+// Call unpatch() to remove the injected behavior.
+export function patch(): void {
   if (unpatchFn !== null) {
     // Don't patch twice.
     return;
   }
 
-  const originalConsoleMethods = { ...targetConsole };
+  const originalConsoleMethods = {};
 
   unpatchFn = () => {
-    for (let method in targetConsole) {
+    for (let method in originalConsoleMethods) {
       try {
         // $FlowFixMe property error|warn is not writable.
         targetConsole[method] = originalConsoleMethods[method];
@@ -61,15 +77,13 @@ export function patch(targetConsole?: Object = console): void {
     }
   };
 
-  for (let method in targetConsole) {
-    const appendComponentStack =
-      method === 'error' || method === 'warn' || method === 'trace';
+  APPEND_STACK_TO_METHODS.forEach(method => {
+    try {
+      const originalMethod = (originalConsoleMethods[method] =
+        targetConsole[method]);
 
-    const originalMethod = targetConsole[method];
-    const overrideMethod = (...args) => {
-      if (isDisabled) return;
-
-      if (appendComponentStack) {
+      // $FlowFixMe property error|warn is not writable.
+      targetConsole[method] = (...args) => {
         // If we are ever called with a string that already has a component stack, e.g. a React error/warning,
         // don't append a second stack.
         const alreadyHasComponentStack =
@@ -105,18 +119,14 @@ export function patch(targetConsole?: Object = console): void {
             }
           }
         }
-      }
 
-      originalMethod(...args);
-    };
-
-    try {
-      // $FlowFixMe property error|warn is not writable.
-      targetConsole[method] = overrideMethod;
+        originalMethod(...args);
+      };
     } catch (error) {}
-  }
+  });
 }
 
+// Removed component stack patch from whitelisted console methods.
 export function unpatch(): void {
   if (unpatchFn !== null) {
     unpatchFn();

src/backend/renderer.js

@@ -40,8 +40,6 @@ import {
 } from '../constants';
 import { inspectHooksOfFiber } from './ReactDebugHooks';
 import {
-  disable as disableConsole,
-  enable as enableConsole,
   patch as patchConsole,
   registerRenderer as registerRendererWithConsole,
 } from './console';
@@ -2254,15 +2252,30 @@ export function attach(
 
     let hooks = null;
     if (usesHooks) {
-      // Suppress console logging while re-rendering
+      const originalConsoleMethods = {};
+
+      // Temporarily disable all console logging before re-running the hook.
+      for (let method in console) {
+        try {
+          originalConsoleMethods[method] = console[method];
+          // $FlowFixMe property error|warn is not writable.
+          console[method] = () => {};
+        } catch (error) {}
+      }
+
       try {
-        disableConsole();
         hooks = inspectHooksOfFiber(
           fiber,
           (renderer.currentDispatcherRef: any)
         );
       } finally {
-        enableConsole();
+        // Restore original console functionality.
+        for (let method in originalConsoleMethods) {
+          try {
+            // $FlowFixMe property error|warn is not writable.
+            console[method] = originalConsoleMethods[method];
+          } catch (error) {}
+        }
       }
     }