From 2e71427d38d41327842552fb45d033fd600f28f6 Mon Sep 17 00:00:00 2001
From: Dylan Kirkby <dsgkirkby@users.noreply.github.com>
Date: Thu, 27 Jul 2017 09:27:16 -0700
Subject: [PATCH] Persist child listeners through hot reload (#715)

---
 src/components/connectAdvanced.js | 17 +++++-
 src/utils/Subscription.js         |  4 ++
 test/components/connect.spec.js   | 91 ++++++++++++++++++++++++++-----
 3 files changed, 95 insertions(+), 17 deletions(-)

diff --git a/src/components/connectAdvanced.js b/src/components/connectAdvanced.js
index 5be9f8cf7..492a0b5da 100644
--- a/src/components/connectAdvanced.js
+++ b/src/components/connectAdvanced.js
@@ -271,9 +271,22 @@ export default function connectAdvanced(
           this.version = version
           this.initSelector()
 
-          if (this.subscription) this.subscription.tryUnsubscribe()
+          // If any connected descendants don't hot reload (and resubscribe in the process), their
+          // listeners will be lost when we unsubscribe. Unfortunately, by copying over all
+          // listeners, this does mean that the old versions of connected descendants will still be
+          // notified of state changes; however, their onStateChange function is a no-op so this
+          // isn't a huge deal.
+          let oldListeners = [];
+
+          if (this.subscription) {
+            oldListeners = this.subscription.listeners.get()
+            this.subscription.tryUnsubscribe()
+          }
           this.initSubscription()
-          if (shouldHandleStateChanges) this.subscription.trySubscribe()
+          if (shouldHandleStateChanges) {
+            this.subscription.trySubscribe()
+            oldListeners.forEach(listener => this.subscription.listeners.subscribe(listener))
+          }
         }
       }
     }
diff --git a/src/utils/Subscription.js b/src/utils/Subscription.js
index 91e96ba09..db8146799 100644
--- a/src/utils/Subscription.js
+++ b/src/utils/Subscription.js
@@ -24,6 +24,10 @@ function createListenerCollection() {
       }
     },
 
+    get() {
+      return next
+    },
+
     subscribe(listener) {
       let isSubscribed = true
       if (next === current) next = current.slice()
diff --git a/test/components/connect.spec.js b/test/components/connect.spec.js
index bdd4725e2..91ab71fb7 100644
--- a/test/components/connect.spec.js
+++ b/test/components/connect.spec.js
@@ -60,6 +60,19 @@ describe('React', () => {
         : prev
     }
 
+    function imitateHotReloading(TargetClass, SourceClass, container) {
+      // Crude imitation of hot reloading that does the job
+      Object.getOwnPropertyNames(SourceClass.prototype).filter(key =>
+        typeof SourceClass.prototype[key] === 'function'
+      ).forEach(key => {
+        if (key !== 'render' && key !== 'constructor') {
+          TargetClass.prototype[key] = SourceClass.prototype[key]
+        }
+      })
+
+      container.forceUpdate()
+    }
+
     it('should receive the store in the context', () => {
       const store = createStore(() => ({}))
 
@@ -1375,28 +1388,76 @@ describe('React', () => {
       expect(stub.props.foo).toEqual(undefined)
       expect(stub.props.scooby).toEqual('doo')
 
-      function imitateHotReloading(TargetClass, SourceClass) {
-        // Crude imitation of hot reloading that does the job
-        Object.getOwnPropertyNames(SourceClass.prototype).filter(key =>
-          typeof SourceClass.prototype[key] === 'function'
-        ).forEach(key => {
-          if (key !== 'render' && key !== 'constructor') {
-            TargetClass.prototype[key] = SourceClass.prototype[key]
-          }
-        })
-
-        container.forceUpdate()
-      }
-
-      imitateHotReloading(ContainerBefore, ContainerAfter)
+      imitateHotReloading(ContainerBefore, ContainerAfter, container)
       expect(stub.props.foo).toEqual('baz')
       expect(stub.props.scooby).toEqual('foo')
 
-      imitateHotReloading(ContainerBefore, ContainerNext)
+      imitateHotReloading(ContainerBefore, ContainerNext, container)
       expect(stub.props.foo).toEqual('bar')
       expect(stub.props.scooby).toEqual('boo')
     })
 
+    it('should persist listeners through hot update', () => {
+      const ACTION_TYPE = "ACTION"
+      const store = createStore((state = {actions: 0}, action) => {
+        switch (action.type) {
+          case ACTION_TYPE: {
+            return {
+              actions: state.actions + 1
+            }
+          }
+          default:
+            return state
+        }
+      })
+
+      @connect(
+        (state) => ({ actions: state.actions })
+      )
+      class Child extends Component {
+        render() {
+          return <Passthrough {...this.props}/>
+        }
+      }
+
+      @connect(
+        () => ({ scooby: 'doo' })
+      )
+      class ParentBefore extends Component {
+        render() {
+          return (
+            <Child />
+          )
+        }
+      }
+
+      @connect(
+        () => ({ scooby: 'boo' })
+      )
+      class ParentAfter extends Component {
+        render() {
+          return (
+            <Child />
+          )
+        }
+      }
+
+      let container
+      TestUtils.renderIntoDocument(
+        <ProviderMock store={store}>
+          <ParentBefore ref={instance => container = instance}/>
+        </ProviderMock>
+      )
+
+      const stub = TestUtils.findRenderedComponentWithType(container, Passthrough)
+
+      imitateHotReloading(ParentBefore, ParentAfter, container)
+
+      store.dispatch({type: ACTION_TYPE})
+
+      expect(stub.props.actions).toEqual(1)
+    })
+
     it('should set the displayName correctly', () => {
       expect(connect(state => state)(
         class Foo extends Component {