Merge branch 'master' of https://github.com/ctrlplusb/react-tree-walker

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(<Baz />, () => true).then(() => {
-      const expected = { foo: 'bar' }
+    return reactTreeWalker(<Baz value="foo" />, () => true).then(() => {
+      const expected = { foo: 'bar', other: 'I am foo bar' }
       expect(actual).toMatchObject(expected)
     })
   })
@@ -164,4 +165,36 @@ describe('reactTreeWalker', () => {
       expect(actual).toMatchObject(expected)
     })
   })
+
+  it('works with instance-as-result component', () => {
+    // eslint-disable-next-line react/prefer-stateless-function
+    class Baz extends Component {
+      render() {
+        return (
+          <div>
+            <Foo something={1} />
+            <Foo something={2} />
+          </div>
+        )
+      }
+    }
+    const Bar = props => new Baz(props)
+    const tree = (
+      <div>
+        <Bar />
+      </div>
+    )
+    const actual = []
+    // eslint-disable-next-line no-unused-vars
+    const visitor = (element, instance, context) => {
+      if (instance && typeof instance.getSomething === 'function') {
+        const something = instance.getSomething()
+        actual.push(something)
+      }
+    }
+    return reactTreeWalker(tree, visitor).then(() => {
+      const expected = [1, 2]
+      expect(actual).toEqual(expected)
+    })
+  })
 })

src/index.js

@@ -49,6 +49,11 @@ const pMapSeries = (iterable, iterator) => {
   ).then(() => ret)
 }
 
+const ensureChild = child =>
+  child && typeof child.render === 'function'
+    ? ensureChild(child.render())
+    : child
+
 export const isPromise = x => x != null && typeof x.then === 'function'
 
 // Recurse an React Element tree, running visitor on each element.
@@ -68,7 +73,7 @@ export default function reactTreeWalker(
           resolve()
         }
 
-        const child = getChildren()
+        const child = ensureChild(getChildren())
         const theChildContext =
           typeof childContext === 'function' ? childContext() : childContext
 
@@ -136,6 +141,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)
         }