Create a jest mock function which is awaitable

waitableJestFn.ts

function waitableJestFn<Result, Args extends any[]>(
  implementation?: (...args: Args) => Result,
) {
  const mock = jest.fn<Result, Args>(implementation);
  const asyncMethods = {
    waitForCallback: (count = 1) =>
      new Promise<Args>((resolve) => {
        mock.mockImplementation(((...args) => {
          count--;
          if (count === 0) resolve(args);
          if (implementation) return implementation(...args);
        }) as (...args: Args) => Result);
      }),
  } as const;
  Object.assign(mock, asyncMethods);
  return mock as typeof mock & typeof asyncMethods;
}