Accessing unique queryClients in Storybook-based Jest tests

tl;dr

I recently came up with a pattern for accessing unique React Query clients on a per-render basis in Storybook stories and tests (using Jest and React Testing Library). This enables the following kind of API:

// composeStories from `@storybook/testing-react`
const Story = composeStories(stories)
const { queryClient } = renderStory(<Story.FancyButton />)
// ...wait for query client state, assert state, etc

(If this sounds like a fun stack to you, we're hiring at Process Street!)

Keep reading for more of the story, or just jump into the code in this template:

I prefer to give each test its own QueryClientProvider and create a new QueryClient for each test. That way, tests are completely isolated from each other. A different approach might be to clear the cache after each test, but I like to keep shared state between tests as minimal as possible. Otherwise, you might get unexpected and flaky results if you run your tests in parallel.

- TkDodo on Testing React Query

...But what if I need access to the query client in my tests!?

At Process Street we use the lovely @storybook/testing-react package to use Storybook stories as our components under test. This is a great pattern because you can effectively automate the manual QA you do in the Storybook browser.

Let's say you just added a toast alert for an API exception case. You finish your code changes and switch to your browser to test the toast in the story. It works! Now you can keep your user hat on and use Testing Library fundamentals to write a Jest spec matching what you did in the browser.

An example (pseudo code):

import * as React from 'react'
// local util module to wrap test utils like React Testing 
// Library (RTL) and @storybook/testing-react
import * as Test from 'test'
import * as stories from './index.stories'

const Story = Test.composeStories(stories)

test("FancyButton shows an alert for failures", () => {
  Test.renderStory(<Story.Failure />)
  Test.user.click(Test.screen.getByText(/do stuff/i))
  await Test.findByText(/uh oh!/i)
})

Yep. That's it. For the past few months our *.spec.tsx files have been very concise and declarative. This is because all the setup is in *.stories.tsx files. Tests just become expressions of how I'm testing the stories, as a user, in the browser.

Along with testing "integrated" components powered by React Query, we've been using a global queryClient instance to leverage patterns like:

await Test.waitFor(() => {
  expect(queryClient.isFetching()).toEq(0)
})

We can't follow TkDodo's advice very easily because if each story sets up it's own query client, how do we get access to it? We could instantiate a new query client for each test, but that's the kind of boilerplate that makes testing feel terrible. My goal is always to make testing feel as good as possible (maybe even fun?!).

The code I wish I had has the following characteristics:

All stories and tests have a unique queryClient instance.
The unique queryClient instance is easily accessible in each test.
The pattern for accessing the queryClient feels like "The Testing Library Way".

By #3, I refer to how Testing Library has normalized our eyes to the pattern of rendering something and destructuring results from that render call.

const { rerender } = Test.render(<FancyButton />)

It would feel awfully nice to be able to do something like:

const { queryClient } = Test.render(<FancyButton />)

Right? It's quite clear that the queryClient is unique to this particular invocation of Test.render.

So the big question is, how do we implement this?

I know right away that I won't instantiate the queryClient nor the QueryClientProvider at the individual story level for the same reasons I wouldn't instantiate it in each test: too much annoying boilerplate that makes writing stories less fun. So that's out. We need some kind of "do this for every test" lever.

It's recommended to have a custom render function that wraps the component under test the same way your app is globally wrapped by some combination of providers. We'll borrow this notion of "all the providers" but skip the custom render wrapper and instead use it for a Storybook decorator. Since we'll want control of our queryClient, we'll parameterize it for a root provider.

// ./test/index.tsx
import React from "react";
import { render, RenderOptions } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "react-query";
import { ChakraProvider } from "@chakra-ui/react";

export const makeQueryClient = () =>
  new QueryClient({
    defaultOptions: { queries: { retry: false } }
  });

type Props = { queryClient?: QueryClient };

export const AllTheProviders: React.FC<Props> = ({
  queryClient = makeQueryClient(),
  children
}) => {
  return (
    <QueryClientProvider client={queryClient}>
      <ChakraProvider>{children}</ChakraProvider>
    </QueryClientProvider>
  );
};

Now we'll jump straight to decorating all stories with AllTheProviders.

// .storybook/main-decorator.tsx
import * as React from "react";
import { AllTheProviders } from "../test";

export const MainDecorator: DecoratorFn = (
  Story,
  options
) => {
  return (
    <AllTheProviders queryClient={options.args.queryClient}>
      <Story {...options} />
    </AllTheProviders>
  );
};

Note that options.args.queryClient is still nullable, but allows us to pass a query client to the component results of composeStories.

Now we just export that decorator for Storybook's browser configuration in preview.js.

// .storybook/preview.js
import { MainDecorator } from './main-decorator'
//...
export const decorators = [AllTheProviders]

Now we have "decorated" stories for testing with composeStories from @storybook/testing-react, but we need a custom render function that adds queryClient to the return value of render from React Testing Library.

export const renderStory = (
  ui: React.ReactElement<{ queryClient?: QueryClient }>,
  options: RenderOptions = {}
) => {
  const queryClient: QueryClient =
    ui.props.queryClient ?? makeQueryClient();
  const clonedUi = React.cloneElement(ui, { queryClient });
  return { ...render(clonedUi, options), queryClient };
};

Boom! We use React.cloneElement to modify the already-invoked component function so we can pass a queryClient from a different scope. If the ui component was already called with a queryClient, that will be reused thanks to our nullish coalescing operator ??. Now in our tests we can access the queryClient as a result of our render call.

const { queryClient } = Test.renderStory(<Story.FancyButton />)

If you do need to test implementation details (which, face it, sometimes you Just Do), you can do something like this:

const queryClient = makeQueryClient()
const invalidateQueriesSpy = 
  jest.spyOn(queryClient, 'invalidateQueries');

Test.render(<Story.Success queryClient={queryClient} />)
Test.user.click(Test.screen.getByText(/do stuff/i))
expect(queryClient.invalidateQueries)
  .toHaveBeenCalledWith("user-profile")