From dbd7f4f820d3cdac8d6d1ff7abca032239851668 Mon Sep 17 00:00:00 2001
From: Kiwi <vinh@wikiwi.io>
Date: Tue, 18 Sep 2018 22:59:51 +0200
Subject: [PATCH] [next] Resubscribe local state (#1884)

* Resubscribe local state

* Don't pass context
---
 .../lib/relay/createMutationContainer.tsx     |  2 +-
 .../lib/relay/withLocalStateContainer.tsx     | 39 ++++++++++++++-----
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/core/client/framework/lib/relay/createMutationContainer.tsx b/src/core/client/framework/lib/relay/createMutationContainer.tsx
index b8c74be18..3084635c9 100644
--- a/src/core/client/framework/lib/relay/createMutationContainer.tsx
+++ b/src/core/client/framework/lib/relay/createMutationContainer.tsx
@@ -42,7 +42,7 @@ function createMutationContainer<T extends string, I, R>(
         };
 
         public render() {
-          const { relayEnvironment: _, ...rest } = this.props;
+          const { context: _, ...rest } = this.props;
           const inject = {
             [propName]: this.commit,
           };
diff --git a/src/core/client/framework/lib/relay/withLocalStateContainer.tsx b/src/core/client/framework/lib/relay/withLocalStateContainer.tsx
index 444b67fa2..484928fc7 100644
--- a/src/core/client/framework/lib/relay/withLocalStateContainer.tsx
+++ b/src/core/client/framework/lib/relay/withLocalStateContainer.tsx
@@ -39,6 +39,7 @@ export const LOCAL_ID = "client:root.local";
 function withLocalStateContainer(
   fragmentSpec: GraphQLTaggedNode
 ): InferableComponentEnhancer<{ local: _RefType<any> }> {
+  const fragment = (fragmentSpec as any).data().default;
   return compose(
     withContext(({ relayEnvironment }) => ({ relayEnvironment })),
     hoistStatics((BaseComponent: React.ComponentType<any>) => {
@@ -49,9 +50,7 @@ function withLocalStateContainer(
         );
         private subscription: Disposable;
 
-        constructor(props: Props) {
-          super(props);
-          const fragment = (fragmentSpec as any).data().default;
+        private subscribe(environment: Environment) {
           if (fragment.kind !== "Fragment") {
             throw new Error("Expected fragment");
           }
@@ -65,24 +64,44 @@ function withLocalStateContainer(
             node: { selections: fragment.selections },
             variables: {},
           };
-          const snapshot = props.relayEnvironment.lookup(selector);
-          this.subscription = props.relayEnvironment.subscribe(
+          const snapshot = environment.lookup(selector);
+          this.subscription = environment.subscribe(
             snapshot,
             this.updateSnapshot
           );
-          this.state = {
-            data: snapshot.data,
-          };
+          this.updateSnapshot(snapshot);
+        }
+
+        constructor(props: Props) {
+          super(props);
+          this.subscribe(props.relayEnvironment);
         }
 
         private updateSnapshot = (snapshot: CSnapshot<any>) => {
-          this.setState({ data: snapshot.data });
+          const nextState = { data: snapshot.data };
+          // State has not been initialized yet.
+          if (!this.state) {
+            this.state = nextState;
+            return;
+          }
+          this.setState(nextState);
         };
 
-        public componentWillUnmount() {
+        private unsubscribe() {
           this.subscription.dispose();
         }
 
+        public componentWillReceiveProps(next: Props) {
+          if (this.props.relayEnvironment !== next.relayEnvironment) {
+            this.unsubscribe();
+            this.subscribe(next.relayEnvironment);
+          }
+        }
+
+        public componentWillUnmount() {
+          this.unsubscribe();
+        }
+
         public render() {
           const { relayEnvironment: _, ...rest } = this.props;
           return <BaseComponent {...rest} local={this.state.data} />;