From 2324174f0c75c995c2540d4e99f451cad7039094 Mon Sep 17 00:00:00 2001
From: zombiej <smith3816@gmail.com>
Date: Mon, 16 Mar 2020 18:38:20 +0800
Subject: [PATCH] fix control treeExpandedKeys

---
 .prettierrc                      |  3 +-
 package.json                     |  1 +
 src/OptionList.tsx               | 53 +++++++++++---------------------
 tests/Select.SearchInput.spec.js | 46 +++++++++++++++++++++++++++
 4 files changed, 67 insertions(+), 36 deletions(-)

diff --git a/.prettierrc b/.prettierrc
index 27dd8afb..f307fb19 100644
--- a/.prettierrc
+++ b/.prettierrc
@@ -3,5 +3,6 @@
   "semi": true,
   "singleQuote": true,
   "tabWidth": 2,
-  "trailingComma": "all"
+  "trailingComma": "all",
+  "printWidth": 100
 }
diff --git a/package.json b/package.json
index c56677cb..27c3399b 100644
--- a/package.json
+++ b/package.json
@@ -44,6 +44,7 @@
     "react-dom": "*"
   },
   "devDependencies": {
+    "@types/jest": "^25.1.4",
     "@types/react": "^16.8.19",
     "@types/react-dom": "^16.8.4",
     "@types/warning": "^3.0.0",
diff --git a/src/OptionList.tsx b/src/OptionList.tsx
index ec578a92..0ff39c87 100644
--- a/src/OptionList.tsx
+++ b/src/OptionList.tsx
@@ -4,13 +4,7 @@ import useMemo from 'rc-util/lib/hooks/useMemo';
 import { RefOptionListProps } from 'rc-select/lib/OptionList';
 import Tree, { TreeProps } from 'rc-tree';
 import { EventDataNode } from 'rc-tree/lib/interface';
-import {
-  FlattenDataNode,
-  RawValueType,
-  DataNode,
-  TreeDataNode,
-  Key,
-} from './interface';
+import { FlattenDataNode, RawValueType, DataNode, TreeDataNode, Key } from './interface';
 import { SelectContext } from './Context';
 import useKeyValueMapping from './hooks/useKeyValueMapping';
 import useKeyValueMap from './hooks/useKeyValueMap';
@@ -56,10 +50,10 @@ export interface OptionListProps<OptionsType extends object[]> {
   onScroll: React.UIEventHandler<HTMLDivElement>;
 }
 
-const OptionList: React.RefForwardingComponent<
-  RefOptionListProps,
-  OptionListProps<DataNode[]>
-> = (props, ref) => {
+const OptionList: React.RefForwardingComponent<RefOptionListProps, OptionListProps<DataNode[]>> = (
+  props,
+  ref,
+) => {
   const {
     prefixCls,
     height,
@@ -102,10 +96,7 @@ const OptionList: React.RefForwardingComponent<
   );
 
   const [cacheKeyMap, cacheValueMap] = useKeyValueMap(flattenOptions);
-  const [getEntityByKey, getEntityByValue] = useKeyValueMapping(
-    cacheKeyMap,
-    cacheValueMap,
-  );
+  const [getEntityByKey, getEntityByValue] = useKeyValueMapping(cacheKeyMap, cacheValueMap);
 
   // ========================== Values ==========================
   const valueKeys = React.useMemo(
@@ -148,14 +139,15 @@ const OptionList: React.RefForwardingComponent<
   };
 
   // =========================== Keys ===========================
-  const [expandedKeys, setExpandedKeys] = React.useState<Key[]>(
-    treeDefaultExpandedKeys,
-  );
-  const [searchExpandedKeys, setSearchExpandedKeys] = React.useState<Key[]>(
-    null,
-  );
-  const mergedExpandedKeys =
-    treeExpandedKeys || (searchValue ? searchExpandedKeys : expandedKeys);
+  const [expandedKeys, setExpandedKeys] = React.useState<Key[]>(treeDefaultExpandedKeys);
+  const [searchExpandedKeys, setSearchExpandedKeys] = React.useState<Key[]>(null);
+
+  const mergedExpandedKeys = React.useMemo(() => {
+    if (treeExpandedKeys) {
+      return [...treeExpandedKeys];
+    }
+    return searchValue ? searchExpandedKeys : expandedKeys;
+  }, [expandedKeys, searchExpandedKeys, lowerSearchValue, treeExpandedKeys]);
 
   React.useEffect(() => {
     if (searchValue) {
@@ -203,9 +195,7 @@ const OptionList: React.RefForwardingComponent<
         case KeyCode.DOWN:
         case KeyCode.LEFT:
         case KeyCode.RIGHT:
-          treeRef.current?.onKeyDown(event as React.KeyboardEvent<
-            HTMLDivElement
-          >);
+          treeRef.current?.onKeyDown(event as React.KeyboardEvent<HTMLDivElement>);
           break;
 
         // >>> Select item
@@ -231,11 +221,7 @@ const OptionList: React.RefForwardingComponent<
   // ========================== Render ==========================
   if (memoOptions.length === 0) {
     return (
-      <div
-        role="listbox"
-        className={`${prefixCls}-empty`}
-        onMouseDown={onListMouseDown}
-      >
+      <div role="listbox" className={`${prefixCls}-empty`} onMouseDown={onListMouseDown}>
         {notFoundContent}
       </div>
     );
@@ -291,10 +277,7 @@ const OptionList: React.RefForwardingComponent<
   );
 };
 
-const RefOptionList = React.forwardRef<
-  RefOptionListProps,
-  OptionListProps<DataNode[]>
->(OptionList);
+const RefOptionList = React.forwardRef<RefOptionListProps, OptionListProps<DataNode[]>>(OptionList);
 RefOptionList.displayName = 'OptionList';
 
 export default RefOptionList;
diff --git a/tests/Select.SearchInput.spec.js b/tests/Select.SearchInput.spec.js
index 89e698c4..4127f184 100644
--- a/tests/Select.SearchInput.spec.js
+++ b/tests/Select.SearchInput.spec.js
@@ -26,4 +26,50 @@ describe('TreeSelect.SearchInput', () => {
         .props().value,
     ).toBeFalsy();
   });
+
+  it('expandedKeys', () => {
+    const wrapper = mount(
+      <TreeSelect
+        open
+        showSearch
+        treeExpandedKeys={['bamboo', 'light']}
+        treeData={[
+          {
+            title: 'bamboo',
+            value: 'bamboo',
+            children: [{ title: '111', value: '111' }],
+          },
+          {
+            title: 'light',
+            value: 'light',
+            children: [{ title: '222', value: '222' }],
+          },
+        ]}
+      />,
+    );
+
+    expect(wrapper.find('NodeList').props().expandedKeys).toEqual(['bamboo', 'light']);
+
+    function search(value) {
+      wrapper
+        .find('input')
+        .first()
+        .simulate('change', { target: { value } });
+      wrapper.update();
+    }
+
+    function listProps() {
+      return wrapper.find('NodeList').props();
+    }
+
+    // Clean up
+    search('bambooA');
+
+    // Return back
+    search('bamboo');
+
+    // Back to default
+    search('');
+    expect(listProps().expandedKeys).toEqual(['bamboo', 'light']);
+  });
 });