From d02d146cfb907d1cb3b2c72a2d59997b32253c16 Mon Sep 17 00:00:00 2001
From: Minko <49073282+hemengke1997@users.noreply.github.com>
Date: Wed, 12 Jun 2024 19:43:39 +0800
Subject: [PATCH] fix(useResetState): reset random initial value incorrectly
 (#2562)

* fix: random initial value changed

* fix(useResetState): functinal initialValue

* chore: revert lock file

* refactor: update

* style: update

---------

Co-authored-by: liuyib <1656081615@qq.com>
---
 .../src/useResetState/__tests__/index.test.ts | 45 +++++++++++++++++-
 .../hooks/src/useResetState/demo/demo1.tsx    | 46 ++++++++++---------
 packages/hooks/src/useResetState/index.ts     | 15 ++++--
 3 files changed, 80 insertions(+), 26 deletions(-)

diff --git a/packages/hooks/src/useResetState/__tests__/index.test.ts b/packages/hooks/src/useResetState/__tests__/index.test.ts
index 8728302445..b4fa76640b 100644
--- a/packages/hooks/src/useResetState/__tests__/index.test.ts
+++ b/packages/hooks/src/useResetState/__tests__/index.test.ts
@@ -2,9 +2,9 @@ import { act, renderHook } from '@testing-library/react';
 import useResetState from '../index';
 
 describe('useResetState', () => {
-  const setUp = <S>(initialState: S) =>
+  const setUp = (initialState) =>
     renderHook(() => {
-      const [state, setState, resetState] = useResetState<S>(initialState);
+      const [state, setState, resetState] = useResetState(initialState);
 
       return {
         state,
@@ -20,6 +20,13 @@ describe('useResetState', () => {
     expect(hook.result.current.state).toEqual({ hello: 'world' });
   });
 
+  it('should support functional initialValue', () => {
+    const hook = setUp(() => ({
+      hello: 'world',
+    }));
+    expect(hook.result.current.state).toEqual({ hello: 'world' });
+  });
+
   it('should reset state', () => {
     const hook = setUp({
       hello: '',
@@ -49,4 +56,38 @@ describe('useResetState', () => {
     });
     expect(hook.result.current.state).toEqual({ count: 1 });
   });
+
+  it('should keep random initial state', () => {
+    const random = Math.random();
+    const hook = setUp({
+      count: random,
+    });
+
+    act(() => {
+      hook.result.current.setState({ count: Math.random() });
+    });
+
+    act(() => {
+      hook.result.current.resetState();
+    });
+
+    expect(hook.result.current.state).toEqual({ count: random });
+  });
+
+  it('should support random functional initialValue', () => {
+    const random = Math.random();
+    const hook = setUp(() => ({
+      count: random,
+    }));
+
+    act(() => {
+      hook.result.current.setState({ count: Math.random() });
+    });
+
+    act(() => {
+      hook.result.current.resetState();
+    });
+
+    expect(hook.result.current.state).toEqual({ count: random });
+  });
 });
diff --git a/packages/hooks/src/useResetState/demo/demo1.tsx b/packages/hooks/src/useResetState/demo/demo1.tsx
index 5a98c7f876..d77bd29ecf 100644
--- a/packages/hooks/src/useResetState/demo/demo1.tsx
+++ b/packages/hooks/src/useResetState/demo/demo1.tsx
@@ -1,33 +1,37 @@
-import React from 'react';
+import React, { useMemo } from 'react';
+import { Button, Space } from 'antd';
 import { useResetState } from 'ahooks';
 
-interface State {
-  hello: string;
-  count: number;
-}
-
 export default () => {
-  const [state, setState, resetState] = useResetState<State>({
+  const initialValue = {
     hello: '',
-    count: 0,
-  });
+    value: Math.random(),
+  };
+  const initialValueMemo = useMemo(() => {
+    return initialValue;
+  }, []);
+
+  const [state, setState, resetState] = useResetState(initialValue);
 
   return (
     <div>
+      <div>initial state: </div>
+      <pre>{JSON.stringify(initialValueMemo, null, 2)}</pre>
+      <div>current state: </div>
       <pre>{JSON.stringify(state, null, 2)}</pre>
-      <p>
-        <button
-          type="button"
-          style={{ marginRight: '8px' }}
-          onClick={() => setState({ hello: 'world', count: 1 })}
+      <Space>
+        <Button
+          onClick={() =>
+            setState(() => ({
+              hello: 'world',
+              value: Math.random(),
+            }))
+          }
         >
-          set hello and count
-        </button>
-
-        <button type="button" onClick={resetState}>
-          resetState
-        </button>
-      </p>
+          set hello and value
+        </Button>
+        <Button onClick={resetState}>resetState</Button>
+      </Space>
     </div>
   );
 };
diff --git a/packages/hooks/src/useResetState/index.ts b/packages/hooks/src/useResetState/index.ts
index bcff86c842..eaa7d6000e 100644
--- a/packages/hooks/src/useResetState/index.ts
+++ b/packages/hooks/src/useResetState/index.ts
@@ -1,16 +1,25 @@
-import { useState } from 'react';
+import { useRef, useState } from 'react';
 import type { Dispatch, SetStateAction } from 'react';
+import { isFunction } from '../utils';
 import useMemoizedFn from '../useMemoizedFn';
+import useCreation from '../useCreation';
 
 type ResetState = () => void;
 
 const useResetState = <S>(
   initialState: S | (() => S),
 ): [S, Dispatch<SetStateAction<S>>, ResetState] => {
-  const [state, setState] = useState(initialState);
+  const initialStateRef = useRef(initialState);
+  const initialStateMemo = useCreation(
+    () =>
+      isFunction(initialStateRef.current) ? initialStateRef.current() : initialStateRef.current,
+    [],
+  );
+
+  const [state, setState] = useState(initialStateMemo);
 
   const resetState = useMemoizedFn(() => {
-    setState(initialState);
+    setState(initialStateMemo);
   });
 
   return [state, setState, resetState];