From af276b28865fb55720aa37c4c3ca7f2e0c5d6878 Mon Sep 17 00:00:00 2001
From: Thomas Florin <thoflo002@gmail.com>
Date: Thu, 18 Aug 2016 23:19:46 -0400
Subject: [PATCH] Adding support to get the component instance that withRouter
 HOC wraps (#3735)

* More code style fixes

* Fixing code style, and updating ref to be a function ref

* updating error message from not passing in proper options object

* Updating withRouter based on comments

* Reverting change to withRouter

* Last minute spacing issues

* Fixing some spacing issues

* withRouter now supports withRef as an option, to better support getting the wrapped instance
---
 modules/__tests__/withRouter-test.js | 18 ++++++++++++++++++
 modules/withRouter.js                | 14 ++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/modules/__tests__/withRouter-test.js b/modules/__tests__/withRouter-test.js
index c083bd74bf..3828c5f508 100644
--- a/modules/__tests__/withRouter-test.js
+++ b/modules/__tests__/withRouter-test.js
@@ -12,6 +12,9 @@ describe('withRouter', function () {
     propTypes: {
       router: routerShape.isRequired
     }
+    testFunction() {
+      return 'hello from the test function'
+    }
     render() {
       expect(this.props.router).toExist()
       return <h1>App</h1>
@@ -59,4 +62,19 @@ describe('withRouter', function () {
 
     render(<Test router={router} test={test} />, node, done)
   })
+
+  it('should support withRefs as a parameter', function (done) {
+    const WrappedApp = withRouter(App, { withRef:true })
+    const router = {
+      push() {},
+      replace() {},
+      go() {},
+      goBack() {},
+      goForward() {},
+      setRouteLeaveHook() {},
+      isActive() {}
+    }
+    const component = render((<WrappedApp router={router}/>), node, done)
+    expect(component.getWrappedInstance().testFunction()).toEqual('hello from the test function')
+  })
 })
diff --git a/modules/withRouter.js b/modules/withRouter.js
index 396c5ad6c9..eccd94e856 100644
--- a/modules/withRouter.js
+++ b/modules/withRouter.js
@@ -1,18 +1,28 @@
 import React from 'react'
 import hoistStatics from 'hoist-non-react-statics'
 import { routerShape } from './PropTypes'
+import warning from './routerWarning'
+
 
 function getDisplayName(WrappedComponent) {
   return WrappedComponent.displayName || WrappedComponent.name || 'Component'
 }
 
-export default function withRouter(WrappedComponent) {
+export default function withRouter(WrappedComponent, options) {
   const WithRouter = React.createClass({
     contextTypes: { router: routerShape },
     propTypes: { router: routerShape },
+    getWrappedInstance() {
+      warning(options && options.withRef, 'To access the wrappedInstance you must provide {withRef : true} as the second argument of the withRouter call')
+      return this._wrappedComponent
+    },
     render() {
       const router = this.props.router || this.context.router
-      return <WrappedComponent {...this.props} router={router} />
+      if (options && options.withRef) {
+        return <WrappedComponent {...this.props} ref={(component)=>this._wrappedComponent = component} router={router} />
+      } else {
+        return <WrappedComponent {...this.props} router={router} />
+      }
     }
   })