fix: TreeSelect removeIcon and clearIcon not working (#18949)

* add tree-select async demo

* TreeSelect支持clearIcon和removeIcon

* remove docs and demo

* run test

* clear code

* add ui test

components/tree-select/__tests__/__snapshots__/index.test.js.snap

@@ -0,0 +1,93 @@
+// Jest Snapshot v1, https://goo.gl/fbAQLP
+
+exports[`TreeSelect TreeSelect Custom Icons should support customized icons 1`] = `
+<span
+  aria-haspopup="listbox"
+  class="ant-select ant-select-enabled ant-select-allow-clear"
+  role="combobox"
+  tabindex="-1"
+>
+  <span
+    class="ant-select-selection ant-select-selection--multiple"
+  >
+    <ul
+      class="ant-select-selection__rendered"
+      role="menubar"
+    >
+      <li
+        class="ant-select-selection__choice"
+        role="menuitem"
+        style="user-select: none;"
+        unselectable="unselectable"
+      >
+        <span
+          class="ant-select-selection__choice__remove"
+        >
+          <span>
+            remove
+          </span>
+        </span>
+        <span
+          class="ant-select-selection__choice__content"
+        >
+          leaf1
+        </span>
+      </li>
+      <li
+        class="ant-select-selection__choice"
+        role="menuitem"
+        style="user-select: none;"
+        unselectable="unselectable"
+      >
+        <span
+          class="ant-select-selection__choice__remove"
+        >
+          <span>
+            remove
+          </span>
+        </span>
+        <span
+          class="ant-select-selection__choice__content"
+        >
+          leaf2
+        </span>
+      </li>
+      <li
+        class="ant-select-search ant-select-search--inline"
+      >
+        <span
+          class="ant-select-search__field__wrap"
+        >
+          <input
+            aria-autocomplete="list"
+            aria-label="filter select"
+            aria-multiline="false"
+            class="ant-select-search__field"
+            style="width: 0px;"
+            type="text"
+            value=""
+          />
+          <span
+            class="ant-select-search__field__mirror"
+          >
+             
+          </span>
+        </span>
+      </li>
+    </ul>
+    <span
+      class="ant-select-selection__clear"
+    >
+      <span>
+        clear
+      </span>
+    </span>
+    <span
+      class="ant-select-search__field__placeholder"
+      style="display: none;"
+    >
+      Please select
+    </span>
+  </span>
+</span>
+`;

components/tree-select/__tests__/index.test.js

@@ -1,6 +1,6 @@
 import React from 'react';
 import { mount } from 'enzyme';
-import TreeSelect from '..';
+import TreeSelect, { TreeNode } from '..';
 import focusTest from '../../../tests/shared/focusTest';
 import mountTest from '../../../tests/shared/mountTest';
 
@@ -17,4 +17,30 @@ describe('TreeSelect', () => {
       expect(multiple.find('.ant-select-search__field').length).toBeTruthy();
     });
   });
+
+  describe('TreeSelect Custom Icons', () => {
+    it('should support customized icons', () => {
+      const wrapper = mount(
+        <TreeSelect
+          showSearch
+          clearIcon={<span>clear</span>}
+          removeIcon={<span>remove</span>}
+          value={['leaf1', 'leaf2']}
+          placeholder="Please select"
+          multiple
+          allowClear
+          treeDefaultExpandAll
+        >
+          <TreeNode value="parent 1" title="parent 1" key="0-1">
+            <TreeNode value="parent 1-0" title="parent 1-0" key="0-1-1">
+              <TreeNode value="leaf1" title="my leaf" key="random" />
+              <TreeNode value="leaf2" title="your leaf" key="random1" />
+            </TreeNode>
+          </TreeNode>
+        </TreeSelect>,
+      );
+
+      expect(wrapper.render()).toMatchSnapshot();
+    });
+  });
 });

components/tree-select/index.tsx

@@ -5,6 +5,7 @@ import omit from 'omit.js';
 import { TreeSelectProps, TreeNodeValue } from './interface';
 import { ConfigConsumer, ConfigConsumerProps } from '../config-provider';
 import warning from '../_util/warning';
+import { cloneElement } from '../_util/reactNode';
 import Icon from '../icon';
 import { AntTreeNodeProps } from '../tree';
 
@@ -74,6 +75,8 @@ export default class TreeSelect<T extends TreeNodeValue> extends React.Component
       dropdownStyle,
       dropdownClassName,
       suffixIcon,
+      removeIcon,
+      clearIcon,
       getPopupContainer,
       ...restProps
     } = this.props;
@@ -99,15 +102,22 @@ export default class TreeSelect<T extends TreeNodeValue> extends React.Component
       checkable = <span className={`${prefixCls}-tree-checkbox-inner`} />;
     }
 
-    const inputIcon = (suffixIcon &&
-      (React.isValidElement<{ className?: string }>(suffixIcon)
-        ? React.cloneElement(suffixIcon)
-        : suffixIcon)) || <Icon type="down" className={`${prefixCls}-arrow-icon`} />;
+    const inputIcon = suffixIcon ? (
+      cloneElement(suffixIcon)
+    ) : (
+      <Icon type="down" className={`${prefixCls}-arrow-icon`} />
+    );
 
-    const removeIcon = <Icon type="close" className={`${prefixCls}-remove-icon`} />;
+    const finalRemoveIcon = removeIcon ? (
+      cloneElement(removeIcon)
+    ) : (
+      <Icon type="close" className={`${prefixCls}-remove-icon`} />
+    );
 
-    const clearIcon = (
-      <Icon type="close-circle" className={`${prefixCls}-clear-icon`} theme="filled" />
+    const finalClearIcon = clearIcon ? (
+      cloneElement(clearIcon)
+    ) : (
+      <Icon type="close-circle" theme="filled" className={`${prefixCls}-clear-icon`} />
     );
 
     return (
@@ -116,8 +126,8 @@ export default class TreeSelect<T extends TreeNodeValue> extends React.Component
           this.renderSwitcherIcon(prefixCls, nodeProps)
         }
         inputIcon={inputIcon}
-        removeIcon={removeIcon}
-        clearIcon={clearIcon}
+        removeIcon={finalRemoveIcon}
+        clearIcon={finalClearIcon}
         {...rest}
         showSearch={showSearch}
         getPopupContainer={getPopupContainer || getContextPopupContainer}

components/tree-select/interface.tsx

@@ -52,6 +52,8 @@ export interface TreeSelectProps<T extends TreeNodeValue> extends AbstractSelect
   searchValue?: string;
   showCheckedStrategy?: 'SHOW_ALL' | 'SHOW_PARENT' | 'SHOW_CHILD';
   suffixIcon?: React.ReactNode;
+  removeIcon?: React.ReactNode;
+  clearIcon?: React.ReactNode;
   treeCheckable?: boolean | React.ReactNode;
   treeCheckStrictly?: boolean;
   treeData?: Array<TreeNode>;