Support passing updater function to setState in SSR mode (apollographql#1263)

* Support setting state via callback in SSR mode

* Add change to Changelog.md

* Support full signature of state updater

Author: mdebbar

Changelog.md

@@ -3,8 +3,8 @@
 ### vNext
 - fix skip on component update.
 - Fix: ensure `client` option can be used with mutation query [#1145](https://github.com/apollographql/react-apollo/pull/1145)
-
 - Made `OptionProps.data`'s `TResult` partial [#1231](https://github.com/apollographql/react-apollo/pull/1231)
+- Support passing an updater function to `setState` in SSR mode [#1263](https://github.com/apollographql/react-apollo/pull/1263)
 - Support arrays being returned from render in SSR [#1158](https://github.com/apollographql/react-apollo/pull/1158)
 
 ### 1.4.16

src/server.ts

@@ -66,6 +66,9 @@ export function walkTree(
       //   however, we don't need to re-render as well only support setState in
       //   componentWillMount, which happens *before* render).
       instance.setState = newState => {
+        if (typeof newState === 'function') {
+          newState = newState(instance.state, instance.props, instance.context);
+        }
         instance.state = assign({}, instance.state, newState);
       };
 

test/react-web/server/index.test.tsx

@@ -646,6 +646,72 @@ describe('SSR', () => {
         .catch(console.error);
     });
 
+    it('should allow for setting state via an updater function', done => {
+      const query = gql`
+        query user($id: ID) {
+          currentUser(id: $id) {
+            firstName
+          }
+        }
+      `;
+      const data = { currentUser: { firstName: 'James' } };
+      const variables = { id: 1 };
+      const networkInterface = mockNetworkInterface({
+        request: { query, variables },
+        result: { data },
+        delay: 50,
+      });
+      const apolloClient = new ApolloClient({
+        networkInterface,
+        addTypename: false,
+      });
+
+      @graphql(query, { name: 'user' })
+      class Element extends React.Component<any, any> {
+        state = {
+          thing: 1,
+          userId: null,
+          client: null,
+        };
+
+        componentWillMount() {
+          this.setState((state, props, context) => ({
+            thing: state.thing + 1,
+            userId: props.id,
+            client: context.client,
+          }));
+        }
+
+        render() {
+          const { user, id } = this.props;
+          expect(this.state.thing).toBe(2);
+          expect(this.state.userId).toBe(id);
+          expect(this.state.client).toBe(apolloClient);
+          return (
+            <div>{user.loading ? 'loading' : user.currentUser.firstName}</div>
+
+          );
+        }
+      }
+
+      const app = (
+        <ApolloProvider client={apolloClient}>
+          <Element id={1} />
+        </ApolloProvider>
+      );
+
+      getDataFromTree(app)
+        .then(() => {
+          const initialState = apolloClient.store.getState();
+          expect(initialState.apollo.data).toBeTruthy();
+          expect(
+            initialState.apollo.data['$ROOT_QUERY.currentUser({"id":1})'],
+          ).toBeTruthy();
+          done();
+        })
+        .catch(console.error);
+    });
+
     it("shouldn't run queries if ssr is turned to off", () => {
       const query = gql`
         query user($id: ID) {