From 0ac8e563d0bdbbbfe12d718710b3959e78d507dd Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <trueadm@users.noreply.github.com>
Date: Fri, 11 Oct 2019 16:58:27 +0200
Subject: [PATCH] [react-interactions] Add getInstanceFromNode support to
 TestHostRenderer (#17065)

Fix bad WeakMap key case

Fix bad WeakMap key case
---
 .../src/__tests__/ReactScope-test.internal.js | 53 +++++++++++++++++++
 .../src/ReactTestHostConfig.js                | 16 ++++--
 2 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js b/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js
index 67c1a088b7818..edf2fbc09621e 100644
--- a/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js
+++ b/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js
@@ -402,6 +402,59 @@ describe('ReactScope', () => {
       expect(node).toEqual(aRef.current);
     });
 
+    it('containsNode() works as intended', () => {
+      const TestScope = React.unstable_createScope((type, props) => true);
+      const scopeRef = React.createRef();
+      const divRef = React.createRef();
+      const spanRef = React.createRef();
+      const aRef = React.createRef();
+      const outerSpan = React.createRef();
+      const emRef = React.createRef();
+
+      function Test({toggle}) {
+        return toggle ? (
+          <div>
+            <span ref={outerSpan}>SPAN</span>
+            <TestScope ref={scopeRef}>
+              <div ref={divRef}>DIV</div>
+              <span ref={spanRef}>SPAN</span>
+              <a ref={aRef}>A</a>
+            </TestScope>
+            <em ref={emRef}>EM</em>
+          </div>
+        ) : (
+          <div>
+            <TestScope ref={scopeRef}>
+              <a ref={aRef}>A</a>
+              <div ref={divRef}>DIV</div>
+              <span ref={spanRef}>SPAN</span>
+              <em ref={emRef}>EM</em>
+            </TestScope>
+            <span ref={outerSpan}>SPAN</span>
+          </div>
+        );
+      }
+
+      const renderer = ReactTestRenderer.create(<Test toggle={true} />, {
+        createNodeMock: element => {
+          return element;
+        },
+      });
+      expect(scopeRef.current.containsNode(divRef.current)).toBe(true);
+      expect(scopeRef.current.containsNode(spanRef.current)).toBe(true);
+      expect(scopeRef.current.containsNode(aRef.current)).toBe(true);
+      expect(scopeRef.current.containsNode(outerSpan.current)).toBe(false);
+      expect(scopeRef.current.containsNode(emRef.current)).toBe(false);
+      renderer.update(<Test toggle={false} />);
+      expect(scopeRef.current.containsNode(divRef.current)).toBe(true);
+      expect(scopeRef.current.containsNode(spanRef.current)).toBe(true);
+      expect(scopeRef.current.containsNode(aRef.current)).toBe(true);
+      expect(scopeRef.current.containsNode(outerSpan.current)).toBe(false);
+      expect(scopeRef.current.containsNode(emRef.current)).toBe(true);
+      renderer.update(<Test toggle={true} />);
+      expect(scopeRef.current.containsNode(emRef.current)).toBe(false);
+    });
+
     it('mixed getParent() and getAllNodes() works as intended', () => {
       const TestScope = React.unstable_createScope((type, props) => true);
       const TestScope2 = React.unstable_createScope((type, props) => true);
diff --git a/packages/react-test-renderer/src/ReactTestHostConfig.js b/packages/react-test-renderer/src/ReactTestHostConfig.js
index 8cbf8bc64918f..ea56e7616aa1b 100644
--- a/packages/react-test-renderer/src/ReactTestHostConfig.js
+++ b/packages/react-test-renderer/src/ReactTestHostConfig.js
@@ -53,6 +53,8 @@ export * from 'shared/HostConfigWithNoHydration';
 const EVENT_COMPONENT_CONTEXT = {};
 const NO_CONTEXT = {};
 const UPDATE_SIGNAL = {};
+const nodeToInstanceMap = new WeakMap();
+
 if (__DEV__) {
   Object.freeze(NO_CONTEXT);
   Object.freeze(UPDATE_SIGNAL);
@@ -62,10 +64,14 @@ export function getPublicInstance(inst: Instance | TextInstance): * {
   switch (inst.tag) {
     case 'INSTANCE':
       const createNodeMock = inst.rootContainerInstance.createNodeMock;
-      return createNodeMock({
+      const mockNode = createNodeMock({
         type: inst.type,
         props: inst.props,
       });
+      if (typeof mockNode === 'object' && mockNode !== null) {
+        nodeToInstanceMap.set(mockNode, inst);
+      }
+      return mockNode;
     default:
       return inst;
   }
@@ -354,6 +360,10 @@ export function unmountFundamentalComponent(
   }
 }
 
-export function getInstanceFromNode(node: Object) {
-  throw new Error('Not yet implemented.');
+export function getInstanceFromNode(mockNode: Object) {
+  const instance = nodeToInstanceMap.get(mockNode);
+  if (instance !== undefined) {
+    return instance.internalInstanceHandle;
+  }
+  return null;
 }