enzyme test of stateful functional component with async state change

App.test.js

import React from 'react';
import { shallow } from 'enzyme';
import App from './App';

it('button click changes color of box', async () => {
  const app = shallow(<App />);
  expect(app.find('.box').length).toEqual(1);
  // cache button element
  const button = app.find('button').last();
  const eventMock = {
    target: {
      getAttribute: function() {
        return button.props()['data-color']
      }
    }
  };
  // pass mocked event object
  await button.props().onClick(eventMock);
  expect(app.find('.box.red').length).toEqual(1);
});