fix: style should not be removed when remounted (#160)

* fix: style should not be removed when remounted * fix: compatible with react 17 * chore: code clean
ant-design · Nov 20, 2023 · a9d7e0b · a9d7e0b
1 parent b8ffb53
commit a9d7e0b
Show file tree

Hide file tree

Showing 3 changed files with 144 additions and 3 deletions.
diff --git a/src/hooks/useGlobalCache.tsx b/src/hooks/useGlobalCache.tsx
@@ -96,7 +96,14 @@ export default function useGlobalCache<CacheType>(
 
           if (nextCount === 0) {
             // Always remove styles in useEffect callback
-            register(() => onCacheRemove?.(cache, false));
+            register(() => {
+              // With polyfill, registered callback will always be called synchronously
+              // But without polyfill, it will be called in effect clean up,
+              // And by that time this cache is cleaned up.
+              if (polyfill || !globalCache.get(fullPath)) {
+                onCacheRemove?.(cache, false);
+              }
+            });
             return null;
           }
 

diff --git a/tests/index.spec.tsx b/tests/index.spec.tsx
@@ -10,6 +10,7 @@ import {
   StyleProvider,
   Theme,
   useCacheToken,
+  useCSSVarRegister,
   useStyleRegister,
 } from '../src';
 import { ATTR_MARK, ATTR_TOKEN, CSS_IN_JS_INSTANCE } from '../src/StyleContext';
@@ -713,4 +714,67 @@ describe('csssinjs', () => {
       });
     });
   });
+
+  describe('should not cleanup style when unmount and mount', () => {
+    const test = (
+      wrapper: (node: ReactElement) => ReactElement = (node) => node,
+    ) => {
+      const spy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+      const Demo = ({
+        myToken,
+        children,
+      }: {
+        myToken?: string;
+        children?: ReactNode;
+      }) => {
+        const [token, hashId] = useCacheToken<DerivativeToken>(
+          theme,
+          [{ primaryColor: myToken }],
+          {
+            salt: 'test',
+          },
+        );
+
+        useCSSVarRegister(
+          {
+            key: 'color',
+            path: ['cssinjs-cleanup-style-when-remount'],
+            token,
+          },
+          () => ({
+            token: token.primaryColor,
+          }),
+        );
+
+        return <div className={classNames('box', hashId)}>{children}</div>;
+      };
+
+      const { rerender } = render(wrapper(<Demo myToken="token1" />));
+      const styles = Array.from(document.head.querySelectorAll('style'));
+      expect(styles).toHaveLength(1);
+
+      rerender(
+        wrapper(
+          <div>
+            <Demo myToken="token1" />
+          </div>,
+        ),
+      );
+      const styles2 = Array.from(document.head.querySelectorAll('style'));
+      expect(styles2).toHaveLength(1);
+
+      spy.mockRestore();
+    };
+
+    it('normal', () => {
+      test();
+    });
+
+    it('strict mode', () => {
+      test((node) => {
+        return <StrictMode>{node}</StrictMode>;
+      });
+    });
+  });
 });
diff --git a/tests/legacy.spec.tsx b/tests/legacy.spec.tsx
@@ -1,9 +1,16 @@
 import { render } from '@testing-library/react';
+import classNames from 'classnames';
+import type { ReactElement, ReactNode } from 'react';
 import * as React from 'react';
-import { useLayoutEffect } from 'react';
+import { StrictMode, useLayoutEffect } from 'react';
 import { expect } from 'vitest';
 import type { CSSInterpolation } from '../src';
-import { Theme, useCacheToken, useStyleRegister } from '../src';
+import {
+  Theme,
+  useCacheToken,
+  useCSSVarRegister,
+  useStyleRegister,
+} from '../src';
 
 interface DesignToken {
   primaryColor: string;
@@ -157,4 +164,67 @@ describe('legacy React version', () => {
     rerender(<Demo show={false} />);
     expect(document.head.querySelectorAll('style')).toHaveLength(2);
   });
+
+  describe('should not cleanup style when unmount and mount', () => {
+    const test = (
+      wrapper: (node: ReactElement) => ReactElement = (node) => node,
+    ) => {
+      const spy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+      const Demo = ({
+        myToken,
+        children,
+      }: {
+        myToken?: string;
+        children?: ReactNode;
+      }) => {
+        const [token, hashId] = useCacheToken<DerivativeToken>(
+          theme,
+          [{ primaryColor: myToken }],
+          {
+            salt: 'test',
+          },
+        );
+
+        useCSSVarRegister(
+          {
+            key: 'color',
+            path: ['cssinjs-cleanup-style-when-remount'],
+            token,
+          },
+          () => ({
+            token: token.primaryColor,
+          }),
+        );
+
+        return <div className={classNames('box', hashId)}>{children}</div>;
+      };
+
+      const { rerender } = render(wrapper(<Demo myToken="token1" />));
+      const styles = Array.from(document.head.querySelectorAll('style'));
+      expect(styles).toHaveLength(1);
+
+      rerender(
+        wrapper(
+          <div>
+            <Demo myToken="token1" />
+          </div>,
+        ),
+      );
+      const styles2 = Array.from(document.head.querySelectorAll('style'));
+      expect(styles2).toHaveLength(1);
+
+      spy.mockRestore();
+    };
+
+    it('normal', () => {
+      test();
+    });
+
+    it('strict mode', () => {
+      test((node) => {
+        return <StrictMode>{node}</StrictMode>;
+      });
+    });
+  });
 });