fix(style-vars): fix css vars on component with suspense as root (#1718)

packages/runtime-dom/__tests__/helpers/useCssVars.spec.ts

@@ -4,7 +4,8 @@ import {
   h,
   reactive,
   nextTick,
-  ComponentOptions
+  ComponentOptions,
+  Suspense
 } from '@vue/runtime-dom'
 
 describe('useCssVars', () => {
@@ -68,6 +69,48 @@ describe('useCssVars', () => {
     }))
   })
 
+  test('on suspense root', async () => {
+    const state = reactive({ color: 'red' })
+    const root = document.createElement('div')
+
+    const AsyncComp = {
+      async setup() {
+        return () => h('p', 'default')
+      }
+    }
+
+    const App = {
+      setup() {
+        useCssVars(() => state)
+        return () =>
+          h(Suspense, null, {
+            default: h(AsyncComp),
+            fallback: h('div', 'fallback')
+          })
+      }
+    }
+
+    render(h(App), root)
+    // css vars use with fallback tree
+    for (const c of [].slice.call(root.children as any)) {
+      expect((c as HTMLElement).style.getPropertyValue(`--color`)).toBe(`red`)
+    }
+    // AsyncComp resolve
+    await nextTick()
+    // Suspense effects flush
+    await nextTick()
+    // css vars use with default tree
+    for (const c of [].slice.call(root.children as any)) {
+      expect((c as HTMLElement).style.getPropertyValue(`--color`)).toBe(`red`)
+    }
+
+    state.color = 'green'
+    await nextTick()
+    for (const c of [].slice.call(root.children as any)) {
+      expect((c as HTMLElement).style.getPropertyValue(`--color`)).toBe('green')
+    }
+  })
+
   test('with <style scoped>', async () => {
     const id = 'data-v-12345'
 

packages/runtime-dom/src/helpers/useCssVars.ts

@@ -39,10 +39,23 @@ function setVarsOnVNode(
   vars: Record<string, string>,
   prefix: string
 ) {
+  if (__FEATURE_SUSPENSE__ && vnode.shapeFlag & ShapeFlags.SUSPENSE) {
+    const { isResolved, isHydrating, fallbackTree, subTree } = vnode.suspense!
+    if (isResolved || isHydrating) {
+      vnode = subTree
+    } else {
+      vnode.suspense!.effects.push(() => {
+        setVarsOnVNode(subTree, vars, prefix)
+      })
+      vnode = fallbackTree
+    }
+  }
+
   // drill down HOCs until it's a non-component vnode
   while (vnode.component) {
     vnode = vnode.component.subTree
   }
+
   if (vnode.shapeFlag & ShapeFlags.ELEMENT && vnode.el) {
     const style = vnode.el.style
     for (const key in vars) {