Add test for legacy context and event loop yielding

Author: kitten

src/__tests__/suspense.test.js

@@ -37,6 +37,41 @@ describe('renderPrepass', () => {
       })
     })
 
+    it('preserves the correct context values across yields', () => {
+      let called = false
+      const Inner = (_, context) => {
+        expect(context.test).toBe(123)
+        called = true
+        return null
+      }
+
+      const Wait = () => {
+        const start = Date.now()
+        while (Date.now() - start < 21) {}
+        return <Inner />
+      }
+
+      class Outer extends Component {
+        getChildContext() {
+          return { test: 123 }
+        }
+
+        render() {
+          return <Wait />
+        }
+      }
+
+      Inner.contextTypes = { test: null }
+      Outer.childContextTypes = { test: null }
+
+      const render$ = renderPrepass(<Outer />)
+      expect(called).toBe(false)
+
+      return render$.then(() => {
+        expect(called).toBe(true)
+      })
+    })
+
     it('does not yields when work is below the threshold', () => {
       const Inner = jest.fn(() => null)
       const Outer = () => <Inner />

src/visitor.js

@@ -190,10 +190,15 @@ const visitLoop = (
       const element = currChildren[currIndex]
       const children = visitElement(element, queue, visitor)
 
-      traversalChildren.push(children)
-      traversalIndex.push(0)
-      traversalMap.push(flushPrevContextMap())
-      traversalStore.push(flushPrevContextStore())
+      if (children.length > 0) {
+        traversalChildren.push(children)
+        traversalIndex.push(0)
+        traversalMap.push(flushPrevContextMap())
+        traversalStore.push(flushPrevContextStore())
+      } else {
+        restoreContextMap(flushPrevContextMap())
+        restoreContextStore(flushPrevContextStore())
+      }
     } else {
       traversalChildren.pop()
       traversalIndex.pop()