Support full signature of state updater

Changelog.md

@@ -5,7 +5,7 @@
 - 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 callback to `setState` in SSR mode [#1263](https://github.com/apollographql/react-apollo/pull/1263)
+- Support passing an updater function to `setState` in SSR mode [#1263](https://github.com/apollographql/react-apollo/pull/1263)
 
 
 ### 1.4.16

src/server.ts

@@ -61,7 +61,7 @@ export function walkTree(
       //   componentWillMount, which happens *before* render).
       instance.setState = newState => {
         if (typeof newState === 'function') {
-          newState = newState(instance.state);
+          newState = newState(instance.state, instance.props, instance.context);
         }
         instance.state = assign({}, instance.state, newState);
       };

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

@@ -615,7 +615,7 @@ describe('SSR', () => {
         .catch(console.error);
     });
 
-    it('should allow for setting state in a component', done => {
+    it.only('should allow for setting state via an updater function', done => {
       const query = gql`
         query user($id: ID) {
           currentUser(id: $id) {
@@ -637,15 +637,25 @@ describe('SSR', () => {
 
       @graphql(query, { name: 'user' })
       class Element extends React.Component<any, any> {
-        state = { thing: 1 };
+        state = {
+          thing: 1,
+          userId: null,
+          client: null,
+        };
 
         componentWillMount() {
-          this.setState(state => ({ thing: state.thing + 1 }));
+          this.setState((state, props, context) => ({
+            thing: state.thing + 1,
+            userId: props.id,
+            client: context.client,
+          }));
         }
 
         render() {
-          const { user } = this.props;
+          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>
           );