✨ handle hooks setState with function

packages/component-test-utils-react/src/__tests__/react-useState.spec.js

@@ -0,0 +1,27 @@
+const {shallow} = require('../shallow');
+const React = require('react');
+
+describe('shallow - react useState', () => {
+  const Component = () => {
+    const [state, setState] = React.useState('initialState');
+    const handleClick = () => {
+      setState(prevState => `${prevState}+1`);
+    };
+
+    return <button type="button" onClick={handleClick}>{state}</button>;
+  };
+
+  it('should return the initialState correctly', () => {
+    const cmp = shallow(<Component/>);
+
+    expect(cmp.html()).toBe('<button type="button" onClick="[handleClick]">initialState</button>');
+  });
+
+  it('should update the state when cllicking on the button', () => {
+    const cmp = shallow(<Component/>);
+
+    cmp.dispatchEvent('click');
+
+    expect(cmp.html()).toBe('<button type="button" onClick="[handleClick]">initialState+1</button>');
+  });
+});

packages/component-test-utils-react/src/dispatcher/index.js

@@ -42,7 +42,12 @@ class Dispatcher {
   }
 
   useState(initialState) {
-    return this.useReducer((_, arg) => arg, initialState);
+    return this.useReducer(
+      (prevState, arg) => {
+        return typeof arg === 'function' ? arg(prevState) : arg;
+      },
+      initialState
+    );
   }
 
   useEffect(fn, memo) {