diff --git a/package.json b/package.json
index 17af047..b977680 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "react-component-tree",
-  "version": "0.2.1",
+  "version": "0.2.2",
   "description": "Serialize and reproduce the state of an entire tree of React components",
   "main": "build/bundle.js",
   "repository": {
diff --git a/src/render.js b/src/render.js
index ce57753..2818a0d 100644
--- a/src/render.js
+++ b/src/render.js
@@ -30,15 +30,15 @@ exports.injectState = function(component, state) {
   var rootState = _.omit(state, 'children'),
       childrenStates = state.children;
 
-  component.setState(rootState);
-
-  if (_.isEmpty(childrenStates)) {
-    return;
-  }
-
-  _.each(component.refs, function(child, ref) {
-    if (!_.isEmpty(childrenStates[ref])) {
-      exports.injectState(child, childrenStates[ref]);
+  component.setState(rootState, function() {
+    if (_.isEmpty(childrenStates)) {
+      return;
     }
+
+    _.each(component.refs, function(child, ref) {
+      if (!_.isEmpty(childrenStates[ref])) {
+        exports.injectState(child, childrenStates[ref]);
+      }
+    });
   });
 };