Apollo MockedProvider wrapper to replace random wait in tests

mockedProvider.jsx

import PropTypes from "prop-types";
import React, { Component, Fragment } from "react";
import { act } from "react-dom/test-utils";
import { useQuery } from "@apollo/react-hooks";
import { MockedProvider as ApolloMockedProvider } from "@apollo/react-testing";


class Deferred {
  constructor() {
    this.resolve = null;
    this.promise = new Promise((resolve) => this.resolve = resolve);
  }

  wait = () => act(() => this.promise);
}


const Hook = ({ mock }) => {
  const { error, data } = useQuery(mock.request.query, mock.request);

  if (!mock.response) {
    const deferred = new Deferred();
    mock.response = deferred.wait;
    mock.response.resolve = deferred.resolve;
  }

  if (error || data) {
    mock.response.resolve(error, data);
  }

  return null;
};


export class MockedProvider extends Component {
  static propTypes = {
    children: PropTypes.any.isRequired,
    addTypename: PropTypes.bool,
    mocks: PropTypes.array,
  }

  static defaultProps = {
    addTypename: false,
    mocks: [],
  }

  render() {
    const { addTypename, mocks, children, ...props } = this.props;

    return (
      <ApolloMockedProvider {...props} addTypename={addTypename} mocks={mocks}>
        <Fragment>
          {mocks.map(this.createHook)}
          {children}
        </Fragment>
      </ApolloMockedProvider>
    );
  }

  createHook = (mock, id) => {
    const { operation } = mock.request.query.definitions.find(({ kind }) => kind === "OperationDefinition");

    if (operation === "query") {
      return <Hook key={id} mock={mock} />;
    }

    return null;
  }
}

myComponent.test.jsx

import React from "react";
import { mount } from "enzyme";

import { MockedProvider } from "./mockedProvider";
import MyComponent from "./index";
import query from "./index.gql";


describe("<MyComponent />", () => {
  let mocks = [];

  beforeEach(() => {
    mocks = [{
      request: { query },
      result: {
        // your data
        data: { },
      },
    }];
  });

  it("renders MyComponent with data", async() => {
    const myComponent = mount(
      <MockedProvider mocks={mocks}>
        <MyComponent />
      </MockedProvider>
    );

    expect(myComponent.isEmptyRender()).toEqual(true);
    // expect to render loading state 

    await mocks[0].response();
    myComponent.update();

    expect(myComponent.isEmptyRender()).toEqual(false);
    // expect to render the data
  });

  it("renders MyComponent with error", async() => {
    mocks[0].error = new Error("Foo bar");
    
    const myComponent = mount(
      <MockedProvider mocks={mocks}>
        <MyComponent />
      </MockedProvider>
    );

    expect(myComponent.isEmptyRender()).toEqual(true);
    // expect to render loading state

    await mocks[0].response();
    myComponent.update();

    expect(myComponent.isEmptyRender()).toEqual(false);
    // expect to render error state
  });
});