From 87379d323debcd52c9ee317a65ca39878d14acbb Mon Sep 17 00:00:00 2001
From: Leland Richardson <leland.m.richardson@gmail.com>
Date: Mon, 14 Mar 2016 22:10:59 -0700
Subject: [PATCH] Fix context method for SFCs

---
 src/ReactWrapper.js           |  3 +++
 src/ReactWrapperComponent.jsx |  2 +-
 test/ReactWrapper-spec.js     | 22 ++++++++++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/ReactWrapper.js b/src/ReactWrapper.js
index 9f1240cd3..57fd827d8 100644
--- a/src/ReactWrapper.js
+++ b/src/ReactWrapper.js
@@ -123,6 +123,9 @@ export default class ReactWrapper {
    * @returns {ReactComponent}
    */
   instance() {
+    if (this.root !== this) {
+      throw new Error('ReactWrapper::instance() can only be called on the root');
+    }
     return this.component.getInstance();
   }
 
diff --git a/src/ReactWrapperComponent.jsx b/src/ReactWrapperComponent.jsx
index f6f640a16..bbe80a355 100644
--- a/src/ReactWrapperComponent.jsx
+++ b/src/ReactWrapperComponent.jsx
@@ -45,7 +45,7 @@ export default function createWrapperComponent(node, options = {}) {
       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 component._instance;
       }
       return inst;
     },
diff --git a/test/ReactWrapper-spec.js b/test/ReactWrapper-spec.js
index 9d125c5b4..90c07ce01 100644
--- a/test/ReactWrapper-spec.js
+++ b/test/ReactWrapper-spec.js
@@ -75,6 +75,28 @@ describeWithDOM('mount', () => {
       expect(wrapper.context().name).to.equal(context.name);
       expect(wrapper.context('name')).to.equal(context.name);
     });
+
+    describeIf(!REACT013, 'stateless components', () => {
+      it('works with stateless components', () => {
+        const Foo = ({ foo }, {_}) => (
+          <div>
+            <div className="bar">bar</div>
+            <div className="qoo">{foo}</div>
+          </div>
+        );
+
+        Foo.contextTypes = {
+          _: React.PropTypes.string,
+        };
+
+        const wrapper = mount(<Foo foo="qux" />, {
+          context: {
+            _: 'foo',
+          },
+        });
+        expect(wrapper.context('_')).to.equal('foo');
+      });
+    });
   });
 
   describeIf(!REACT013, 'stateless components', () => {