From 184af0b56f079c381fb12d1bc35890efe6601800 Mon Sep 17 00:00:00 2001
From: jeremywu <15975785+JeremyWuuuuu@users.noreply.github.com>
Date: Tue, 16 Mar 2021 21:40:45 +0800
Subject: [PATCH] fix(hooks): fix use-lock-screen hook (#1651)

- Fix when unmounting `use-lock-screen` could cleanup the side effects
---
 .../hooks/__tests__/use-lockscreen.spec.ts    | 61 ++++++++++++++++---
 packages/hooks/use-lockscreen/index.ts        | 19 ++++--
 2 files changed, 67 insertions(+), 13 deletions(-)

diff --git a/packages/hooks/__tests__/use-lockscreen.spec.ts b/packages/hooks/__tests__/use-lockscreen.spec.ts
index 3ce02dbe68..064d7d8abd 100644
--- a/packages/hooks/__tests__/use-lockscreen.spec.ts
+++ b/packages/hooks/__tests__/use-lockscreen.spec.ts
@@ -1,15 +1,62 @@
-import { ref, nextTick } from 'vue'
+import { ref, nextTick, defineComponent, onMounted } from 'vue'
+import { mount } from '@vue/test-utils'
 import { hasClass } from '@element-plus/utils/dom'
+
 import useLockScreen from '../use-lockscreen'
 
+const kls = 'el-popup-parent--hidden'
+
+const Comp = defineComponent({
+  setup() {
+    const flag = ref(false)
+    useLockScreen(flag)
+    onMounted(() => {
+      flag.value = true
+    })
+  },
+  template: `<div></div>`,
+})
+
 describe('useLockScreen', () => {
   test('should lock screen when trigger is true', async () => {
-    const _ref = ref(false)
-    const cls = 'el-popup-parent--hidden'
-    useLockScreen(_ref)
-    expect(hasClass(document.body, cls)).toBe(false)
-    _ref.value = true
+    const wrapper = mount({
+      template: `
+        <test-comp />
+      `,
+      components: {
+        'test-comp': Comp,
+      },
+    })
     await nextTick()
-    expect(hasClass(document.body, cls)).toBe(true)
+    expect(hasClass(document.body, kls)).toBe(true)
+
+    wrapper.unmount()
+    await nextTick()
+    expect(hasClass(document.body, kls)).toBe(false)
+  })
+
+  test('should cleanup when unmounted', async () => {
+    const wrapper = mount({
+      template: `
+        <test-comp v-if="shouldRender" />
+      `,
+      data() {
+        return {
+          shouldRender: true,
+        }
+      },
+      components: {
+        'test-comp': Comp,
+      },
+    })
+
+    await nextTick()
+
+    expect(hasClass(document.body, kls)).toBe(true)
+
+    wrapper.vm.shouldRender = false
+    await nextTick()
+
+    expect(hasClass(document.body, kls)).toBe(false)
   })
 })
diff --git a/packages/hooks/use-lockscreen/index.ts b/packages/hooks/use-lockscreen/index.ts
index 54a91db5dd..23c2c9803a 100644
--- a/packages/hooks/use-lockscreen/index.ts
+++ b/packages/hooks/use-lockscreen/index.ts
@@ -1,4 +1,4 @@
-import { watch, isRef } from 'vue'
+import { watch, isRef, onUnmounted } from 'vue'
 
 import getScrollBarWidth from '@element-plus/utils/scrollbar-width'
 import throwError from '@element-plus/utils/error'
@@ -27,6 +27,17 @@ export default (trigger: Ref<boolean>) => {
   let withoutHiddenClass = false
   let bodyPaddingRight = '0'
   let computedBodyPaddingRight = 0
+
+  onUnmounted(() => {
+    cleanup()
+  })
+
+  const cleanup = () => {
+    removeClass(document.body, 'el-popup-parent--hidden')
+    if (withoutHiddenClass) {
+      document.body.style.paddingRight = bodyPaddingRight
+    }
+  }
   watch(trigger, val => {
     if (val) {
       withoutHiddenClass = !hasClass(document.body, 'el-popup-parent--hidden')
@@ -51,11 +62,7 @@ export default (trigger: Ref<boolean>) => {
       }
       addClass(document.body, 'el-popup-parent--hidden')
     } else {
-      if (withoutHiddenClass) {
-        document.body.style.paddingRight = bodyPaddingRight
-        removeClass(document.body, 'el-popup-parent--hidden')
-      }
-      withoutHiddenClass = true
+      cleanup()
     }
   })
 }