Testing React State with Hooks, Jest, and Enzyme

FunctionalComponent.jsx

import React from 'react';

const TestComponent = () => {
  const [count, setCount] = React.useState(0);

  return (
    <h3>{count}</h3>
    <span>
      <button id="count-up" type="button" onClick={() => setCount(count + 1)}>Count Up</button>
      <button id="count-down" type="button" onClick={() => setCount(count - 1)}>Count Down</button>
      <button id="zero-count" type="button" onClick={() => setCount(0)}>Zero</button>
    </span>
  );
}

export default TestComponent;

FunctionalComponent.test.jsx

import React from 'react';
import Enzyme from 'enzyme';
import Adapter from 'enzyme-adapter-react-16';
import TestComponent from './FunctionalComponent';

Enzyme.configure({ adapter: new Adapter() });

describe('<TestComponent />', () => {
  let wrapper;
  const setState = jest.fn();
  const useStateSpy = jest.spyOn(React, 'useState')
  useStateSpy.mockImplementation((init) => [init, setState]);

  beforeEach(() => {
    wrapper = Enzyme.shallow(<TestComponent />);
  });

  afterEach(() => {
    jest.clearAllMocks();
  });

  describe('Count Up', () => {
    it('calls setCount with count + 1', () => {
      wrapper.find('#count-up').props().onClick();
      expect(setState).toHaveBeenCalledWith(1);
    });
  });

  describe('Count Down', () => {
    it('calls setCount with count - 1', () => {
      wrapper.find('#count-down').props().onClick();
      expect(setState).toHaveBeenCalledWith(-1);
    });
  });

  describe('Zero', () => {
    it('calls setCount with 0', () => {
      wrapper.find('#zero-count').props().onClick();
      expect(setState).toHaveBeenCalledWith(0);
    });
  });
});