diff --git a/src/index.js b/src/index.js
index cda36e9..8f54493 100644
--- a/src/index.js
+++ b/src/index.js
@@ -216,7 +216,7 @@ export default function reactTreeWalker(
               instance.state = instance.state || null
 
               // Make the setState synchronous.
-              instance.setState = newState => {
+              instance.setState = (newState, cb) => {
                 if (typeof newState === 'function') {
                   // eslint-disable-next-line no-param-reassign
                   newState = newState(
@@ -226,6 +226,9 @@ export default function reactTreeWalker(
                   )
                 }
                 instance.state = Object.assign({}, instance.state, newState)
+                if (typeof cb === 'function'){
+                  cb();
+                }
               }
 
               if (Component.getDerivedStateFromProps) {