From e9273abe3d08a9d7a0f0d18b9341d4f22d361df2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=BA=8C=E8=B4=A7=E7=88=B1=E5=90=83=E7=99=BD=E8=90=9D?=
 =?UTF-8?q?=E5=8D=9C?= <smith3816@gmail.com>
Date: Fri, 20 Jan 2023 21:06:13 +0800
Subject: [PATCH] fix: Badge `findDOMNode` warning (#40347)

* fix: findDOMNode warning

* test: add test case
---
 components/badge/ScrollNumber.tsx         | 120 ++++++++++++----------
 components/badge/__tests__/index.test.tsx |  20 +++-
 components/badge/index.tsx                |   3 +-
 3 files changed, 88 insertions(+), 55 deletions(-)

diff --git a/components/badge/ScrollNumber.tsx b/components/badge/ScrollNumber.tsx
index f246287fea..290dd1023a 100644
--- a/components/badge/ScrollNumber.tsx
+++ b/components/badge/ScrollNumber.tsx
@@ -10,7 +10,7 @@ export interface ScrollNumberProps {
   motionClassName?: string;
   count?: string | number | null;
   children?: React.ReactElement<HTMLElement>;
-  component?: string;
+  component?: React.ComponentType<any>;
   style?: React.CSSProperties;
   title?: string | number | null;
   show: boolean;
@@ -21,61 +21,75 @@ export interface ScrollNumberState {
   count?: string | number | null;
 }
 
-const ScrollNumber: React.FC<ScrollNumberProps> = ({
-  prefixCls: customizePrefixCls,
-  count,
-  className,
-  motionClassName,
-  style,
-  title,
-  show,
-  component = 'sup',
-  children,
-  ...restProps
-}) => {
-  const { getPrefixCls } = React.useContext(ConfigContext);
-  const prefixCls = getPrefixCls('scroll-number', customizePrefixCls);
+const ScrollNumber = React.forwardRef<React.Ref<HTMLDivElement>, ScrollNumberProps>(
+  (
+    {
+      prefixCls: customizePrefixCls,
+      count,
+      className,
+      motionClassName,
+      style,
+      title,
+      show,
+      component: Component = 'sup',
+      children,
+      ...restProps
+    },
+    ref,
+  ) => {
+    const { getPrefixCls } = React.useContext(ConfigContext);
+    const prefixCls = getPrefixCls('scroll-number', customizePrefixCls);
 
-  // ============================ Render ============================
-  const newProps = {
-    ...restProps,
-    'data-show': show,
-    style,
-    className: classNames(prefixCls, className, motionClassName),
-    title: title as string,
-  };
+    // ============================ Render ============================
+    const newProps = {
+      ...restProps,
+      'data-show': show,
+      style,
+      className: classNames(prefixCls, className, motionClassName),
+      title: title as string,
+    };
 
-  // Only integer need motion
-  let numberNodes: React.ReactNode = count;
-  if (count && Number(count) % 1 === 0) {
-    const numberList = String(count).split('');
+    // Only integer need motion
+    let numberNodes: React.ReactNode = count;
+    if (count && Number(count) % 1 === 0) {
+      const numberList = String(count).split('');
 
-    numberNodes = numberList.map((num, i) => (
-      <SingleNumber
-        prefixCls={prefixCls}
-        count={Number(count)}
-        value={num}
-        // eslint-disable-next-line react/no-array-index-key
-        key={numberList.length - i}
-      />
-    ));
-  }
+      numberNodes = numberList.map((num, i) => (
+        <SingleNumber
+          prefixCls={prefixCls}
+          count={Number(count)}
+          value={num}
+          // eslint-disable-next-line react/no-array-index-key
+          key={numberList.length - i}
+        />
+      ));
+    }
 
-  // allow specify the border
-  // mock border-color by box-shadow for compatible with old usage:
-  // <Badge count={4} style={{ backgroundColor: '#fff', color: '#999', borderColor: '#d9d9d9' }} />
-  if (style && style.borderColor) {
-    newProps.style = {
-      ...style,
-      boxShadow: `0 0 0 1px ${style.borderColor} inset`,
-    };
-  }
-  if (children) {
-    return cloneElement(children, (oriProps) => ({
-      className: classNames(`${prefixCls}-custom-component`, oriProps?.className, motionClassName),
-    }));
-  }
-  return React.createElement(component, newProps, numberNodes);
-};
+    // allow specify the border
+    // mock border-color by box-shadow for compatible with old usage:
+    // <Badge count={4} style={{ backgroundColor: '#fff', color: '#999', borderColor: '#d9d9d9' }} />
+    if (style && style.borderColor) {
+      newProps.style = {
+        ...style,
+        boxShadow: `0 0 0 1px ${style.borderColor} inset`,
+      };
+    }
+    if (children) {
+      return cloneElement(children, (oriProps) => ({
+        className: classNames(
+          `${prefixCls}-custom-component`,
+          oriProps?.className,
+          motionClassName,
+        ),
+      }));
+    }
+
+    return (
+      <Component {...newProps} ref={ref}>
+        {numberNodes}
+      </Component>
+    );
+  },
+);
 
 export default ScrollNumber;
diff --git a/components/badge/__tests__/index.test.tsx b/components/badge/__tests__/index.test.tsx
index 14544bc6b5..73a680dee9 100644
--- a/components/badge/__tests__/index.test.tsx
+++ b/components/badge/__tests__/index.test.tsx
@@ -25,6 +25,24 @@ describe('Badge', () => {
     jest.useRealTimers();
   });
 
+  it('no strict warning', () => {
+    const errSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+    const { rerender } = render(
+      <Badge dot>
+        <span />
+      </Badge>,
+    );
+
+    rerender(
+      <Badge>
+        <span />
+      </Badge>,
+    );
+
+    expect(errSpy).not.toHaveBeenCalled();
+    errSpy.mockRestore();
+  });
+
   it('badge dot not scaling count > 9', () => {
     const { container } = render(<Badge count={10} dot />);
     expect(container.querySelectorAll('.ant-card-multiple-words').length).toBe(0);
@@ -193,7 +211,7 @@ describe('Badge', () => {
         <Badge count={0} showZero color="#ff0" />
         <Badge count={0} showZero color="blue" />
         <Badge count={0} showZero />
-        <Badge count={0} showZero color='green'>
+        <Badge count={0} showZero color="green">
           <div />
         </Badge>
       </>,
diff --git a/components/badge/index.tsx b/components/badge/index.tsx
index 22d7732660..0a227c6ad0 100644
--- a/components/badge/index.tsx
+++ b/components/badge/index.tsx
@@ -196,7 +196,7 @@ const Badge: CompoundedComponent = ({
         motionAppear={false}
         motionDeadline={1000}
       >
-        {({ className: motionClassName }) => {
+        {({ className: motionClassName, ref }) => {
           const scrollNumberPrefixCls = getPrefixCls(
             'scroll-number',
             customizeScrollNumberPrefixCls,
@@ -230,6 +230,7 @@ const Badge: CompoundedComponent = ({
               title={titleNode}
               style={scrollNumberStyle}
               key="scrollNumber"
+              ref={ref}
             >
               {displayNode}
             </ScrollNumber>