[next] Resubscribe local state (#1884)

* Resubscribe local state

* Don't pass context

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,
           };

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} />;