diff --git a/src/__tests__/index.test.js b/src/__tests__/index.test.js
index 57b1910..c0df628 100644
--- a/src/__tests__/index.test.js
+++ b/src/__tests__/index.test.js
@@ -101,6 +101,7 @@ describe('reactTreeWalker', () => {
componentWillMount() {
this.setState({ foo: 'bar' })
+ this.setState((state, props) => ({ other: `I am ${props.value} ${state.foo}` }))
}
render() {
@@ -109,8 +110,8 @@ describe('reactTreeWalker', () => {
}
}
- return reactTreeWalker(, () => true).then(() => {
- const expected = { foo: 'bar' }
+ return reactTreeWalker(, () => true).then(() => {
+ const expected = { foo: 'bar', other: 'I am foo bar' }
expect(actual).toMatchObject(expected)
})
})
diff --git a/src/index.js b/src/index.js
index 1144205..cfd7696 100644
--- a/src/index.js
+++ b/src/index.js
@@ -136,6 +136,14 @@ export default function reactTreeWalker(
// Make the setState synchronous.
instance.setState = newState => {
+ if (typeof newState === 'function') {
+ // eslint-disable-next-line no-param-reassign
+ newState = newState(
+ instance.state,
+ instance.props,
+ instance.context,
+ )
+ }
instance.state = Object.assign({}, instance.state, newState)
}