From 1a43f2701d2f0eddc5af2af0bf3e26a45874a305 Mon Sep 17 00:00:00 2001
From: Ian Christian Myers <ian@iancmyers.com>
Date: Sat, 19 Mar 2016 17:20:27 -0700
Subject: [PATCH] Fix issue finding nested SFCs by constructor in mount

This fixes an issue where calling `.find` with an SFC constructor/
function would fail to return instances of the the SFC in the tree.

Fixes #262
---
 src/MountedTraversal.js   | 11 ++++++++---
 src/Utils.js              |  4 ++++
 test/ReactWrapper-spec.js |  2 ++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/MountedTraversal.js b/src/MountedTraversal.js
index bf3343eb5..ad1b94326 100644
--- a/src/MountedTraversal.js
+++ b/src/MountedTraversal.js
@@ -5,6 +5,7 @@ import {
   coercePropValue,
   nodeEqual,
   propsOfNode,
+  isFunctionalComponent,
   isSimpleSelector,
   splitSelector,
   selectorError,
@@ -60,12 +61,17 @@ export function instHasId(inst, id) {
   return instId === id;
 }
 
+function isFunctionalComponentWithType(inst, func) {
+  return isFunctionalComponent(inst) && getNode(inst).type === func;
+}
+
 export function instHasType(inst, type) {
   switch (typeof type) {
     case 'string':
       return nodeHasType(getNode(inst), type);
     case 'function':
-      return isCompositeComponentWithType(inst, type);
+      return isCompositeComponentWithType(inst, type) ||
+        isFunctionalComponentWithType(inst, type);
     default:
       return false;
   }
@@ -249,8 +255,7 @@ function findAllInRenderedTreeInternal(inst, test) {
     const internal = internalInstance(inst);
     return findAllInRenderedTreeInternal(internal, test);
   }
-
-  const publicInst = inst.getPublicInstance();
+  const publicInst = inst.getPublicInstance() || inst._instance;
   let ret = test(publicInst) ? [publicInst] : [];
   const currentElement = inst._currentElement;
   if (isDOMComponent(publicInst)) {
diff --git a/src/Utils.js b/src/Utils.js
index a02156c68..e5d017932 100644
--- a/src/Utils.js
+++ b/src/Utils.js
@@ -19,6 +19,10 @@ export function internalInstance(inst) {
     inst[internalInstanceKey(inst)];
 }
 
+export function isFunctionalComponent(inst) {
+  return inst && inst.constructor && inst.constructor.name === 'StatelessComponent';
+}
+
 export function propsOfNode(node) {
   if (REACT013 && node && node._store) {
     return (node._store.props) || {};
diff --git a/test/ReactWrapper-spec.js b/test/ReactWrapper-spec.js
index ebab5fced..cb22b90b9 100644
--- a/test/ReactWrapper-spec.js
+++ b/test/ReactWrapper-spec.js
@@ -128,6 +128,8 @@ describeWithDOM('mount', () => {
       const children = wrapper.find('.box').children();
       expect(children).to.have.length(3);
       expect(children.at(0).props().test).to.equal('123');
+      expect(wrapper.find(TestItem)).to.have.length(3);
+      expect(wrapper.find(TestItem).first().props().test).to.equal('123');
     })
   });