From 274731a784cbd411bc59015155a98f3c8323bab7 Mon Sep 17 00:00:00 2001
From: Leland Richardson <leland.m.richardson@gmail.com>
Date: Sat, 5 Dec 2015 10:11:51 -0800
Subject: [PATCH] Support for Stateless Components

---
 src/ReactWrapper.js                  |  4 ++--
 src/ReactWrapperComponent.jsx        | 22 +++++++++++++++++++++-
 src/__tests__/ReactWrapper-spec.js   | 17 +++++++++++++++++
 src/__tests__/ShallowWrapper-spec.js | 18 +++++++++++++++++-
 src/__tests__/_helpers.js            | 11 +++++++++++
 5 files changed, 68 insertions(+), 4 deletions(-)
 create mode 100644 src/__tests__/_helpers.js

diff --git a/src/ReactWrapper.js b/src/ReactWrapper.js
index 7d34d3b57..4eb86dab2 100644
--- a/src/ReactWrapper.js
+++ b/src/ReactWrapper.js
@@ -61,7 +61,7 @@ export default class ReactWrapper {
           />
       );
       this.root = this;
-      this.node = this.component.refs.component;
+      this.node = this.component.getWrappedComponent();
       this.nodes = [this.node];
       this.length = 1;
     } else {
@@ -108,7 +108,7 @@ export default class ReactWrapper {
    * @returns {ReactComponent}
    */
   instance() {
-    return this.component.refs.component;
+    return this.component.getInstance();
   }
 
   /**
diff --git a/src/ReactWrapperComponent.jsx b/src/ReactWrapperComponent.jsx
index 50a75dbc9..468dce8df 100644
--- a/src/ReactWrapperComponent.jsx
+++ b/src/ReactWrapperComponent.jsx
@@ -19,10 +19,30 @@ export default class ReactWrapperComponent extends React.Component {
     return new Promise(resolve => this.setState(newProps, resolve));
   }
 
+  getInstance() {
+    const component = this._reactInternalInstance._renderedComponent;
+    const inst = component.getPublicInstance();
+    if (inst === null) {
+      throw new Error(
+        `You cannot get an instance of a stateless component.`
+      );
+    }
+    return inst;
+  }
+
+  getWrappedComponent() {
+    const component = this._reactInternalInstance._renderedComponent;
+    const inst = component.getPublicInstance();
+    if (inst === null) {
+      return component;
+    }
+    return inst;
+  }
+
   render() {
     const { Component } = this.props;
     return (
-      <Component ref="component" {...this.state} />
+      <Component {...this.state} />
     );
   }
 }
diff --git a/src/__tests__/ReactWrapper-spec.js b/src/__tests__/ReactWrapper-spec.js
index ca065d699..c6ce25398 100644
--- a/src/__tests__/ReactWrapper-spec.js
+++ b/src/__tests__/ReactWrapper-spec.js
@@ -7,9 +7,26 @@ import {
   ReactWrapper,
   describeWithDOM,
 } from '../';
+import { describeIf } from './_helpers';
+import { REACT013 } from '../version';
 
 describeWithDOM('mount', () => {
 
+  describeIf(!REACT013, 'stateless components', () => {
+    it('works with stateless components', () => {
+      const Foo = ({ foo }) => (
+        <div>
+          <div className="bar">bar</div>
+          <div className="qoo">{foo}</div>
+        </div>
+      );
+      const wrapper = mount(<Foo foo="qux" />);
+      expect(wrapper.type()).to.equal(Foo);
+      expect(wrapper.find('.bar')).to.have.length(1);
+      expect(wrapper.find('.qoo').text()).to.equal('qux');
+    });
+  });
+
   describe('.contains(node)', () => {
 
     it('should allow matches on the root node', () => {
diff --git a/src/__tests__/ShallowWrapper-spec.js b/src/__tests__/ShallowWrapper-spec.js
index eef01c980..ba98aac85 100644
--- a/src/__tests__/ShallowWrapper-spec.js
+++ b/src/__tests__/ShallowWrapper-spec.js
@@ -2,10 +2,26 @@ import React from 'react';
 import { expect } from 'chai';
 import { shallow, render, ShallowWrapper } from '../';
 import sinon from 'sinon';
-
+import { describeIf } from './_helpers';
+import { REACT013 } from '../version';
 
 describe('shallow', () => {
 
+  describeIf(!REACT013, 'stateless components', () => {
+    it('works with stateless components', () => {
+      const Foo = ({ foo }) => (
+        <div>
+          <div className="bar">bar</div>
+          <div className="qoo">{foo}</div>
+        </div>
+      );
+      const wrapper = shallow(<Foo foo="qux" />);
+      expect(wrapper.type()).to.equal('div');
+      expect(wrapper.find('.bar')).to.have.length(1);
+      expect(wrapper.find('.qoo').text()).to.equal('qux');
+    });
+  });
+
   describe('.contains(node)', () => {
 
     it('should allow matches on the root node', () => {
diff --git a/src/__tests__/_helpers.js b/src/__tests__/_helpers.js
new file mode 100644
index 000000000..e98469969
--- /dev/null
+++ b/src/__tests__/_helpers.js
@@ -0,0 +1,11 @@
+/**
+ * Simple wrapper around mocha describe which allows a boolean to be passed in first which
+ * determines whether or not the test will be run
+ */
+export function describeIf(test, a, b) {
+  if (test) {
+    describe(a, b);
+  } else {
+    describe.skip(a, b);
+  }
+}