From 202272488ab76a30405bd103474d9912313ea122 Mon Sep 17 00:00:00 2001
From: FBerthelot <florent.berthelot35@gmail.com>
Date: Thu, 19 Dec 2019 13:35:44 +0100
Subject: [PATCH] :sparkles: handle hooks setState with function

---
 .../src/__tests__/react-useState.spec.js      | 27 +++++++++++++++++++
 .../src/dispatcher/index.js                   |  7 ++++-
 2 files changed, 33 insertions(+), 1 deletion(-)
 create mode 100644 packages/component-test-utils-react/src/__tests__/react-useState.spec.js

diff --git a/packages/component-test-utils-react/src/__tests__/react-useState.spec.js b/packages/component-test-utils-react/src/__tests__/react-useState.spec.js
new file mode 100644
index 0000000..6bbd558
--- /dev/null
+++ b/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>');
+  });
+});
diff --git a/packages/component-test-utils-react/src/dispatcher/index.js b/packages/component-test-utils-react/src/dispatcher/index.js
index 6688e01..75c3815 100644
--- a/packages/component-test-utils-react/src/dispatcher/index.js
+++ b/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) {