From 07ece2e9260fe30a50e7cf317d2ff69f113ecad1 Mon Sep 17 00:00:00 2001
From: underfin <2218301630@qq.com>
Date: Wed, 29 Jul 2020 03:30:18 +0800
Subject: [PATCH] fix(style-vars): fix css vars on component with suspense as
 root (#1718)

---
 .../__tests__/helpers/useCssVars.spec.ts      | 45 ++++++++++++++++++-
 .../runtime-dom/src/helpers/useCssVars.ts     | 13 ++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/packages/runtime-dom/__tests__/helpers/useCssVars.spec.ts b/packages/runtime-dom/__tests__/helpers/useCssVars.spec.ts
index 5260d2a82..293dc7049 100644
--- a/packages/runtime-dom/__tests__/helpers/useCssVars.spec.ts
+++ b/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'
 
diff --git a/packages/runtime-dom/src/helpers/useCssVars.ts b/packages/runtime-dom/src/helpers/useCssVars.ts
index 42f5d6625..783eaaea3 100644
--- a/packages/runtime-dom/src/helpers/useCssVars.ts
+++ b/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) {