Zustand store pattern proposal

[gmaclennan]

This is a proposal for a pattern for creating Zustand stores. The goals of this proposed pattern are:

1. Avoid creating global singletons of store instances, which results in complex test setup to clear the store between each test (see Zustand testing docs)
2. Expose store instances to enable integrations outside of the react render cycle
3. Keep actions out of the Zustand store, to avoid possible mutations of actions and challenges when persisting.

Some ideas for this pattern come from https://tkdodo.eu/blog/zustand-and-react-context

Each "store" consists of a Zustand store API, which only contains the state, and "actions", which is a map of functions to mutate the state.

Better naming would help distinguish between a "Zustand store" which we will use only for state, and our "store" which contains both the Zustand store and actions.

Each "store" should export:

1. A store creation function
2. A provider, which takes a store instance as a value prop (could rename this default)
3. A useStoreState(selector) hook which returns the state of the store
4. A useStoreActions() hook which returns actions to mutate the store

E.g.

import {createContext, useContext} from 'react';
import {createStore, useStore, type StoreApi} from 'zustand';
import {persist as createPersistedState} from 'zustand/middleware';

type BearState = {
  bears: number;
};

function createInitialState(): BearState {
  return {
    bears: 0,
  };
}

const BEAR_STORE_PERSIST_KEY = 'bearStore';

export function createBearStore({persist = false}: {persist: boolean}) {
  let store: StoreApi<BearState>;

  if (persist) {
    store = createStore(
      createPersistedState(createInitialState, {name: BEAR_STORE_PERSIST_KEY}),
    );
  } else {
    store = createStore(createInitialState);
  }

  const actions = {
    increasePopulation() {
      store.setState(state => ({bears: state.bears + 1}));
    },
    removeAllBears: () => store.setState({bears: 0}),
    updateBears: (newBears: number) => store.setState({bears: newBears}),
  };

  return {store, actions};
}

const BearStoreContext = createContext<ReturnType<
  typeof createBearStore
> | null>(null);

export const BearStoreProvider = BearStoreContext.Provider;

// Shouldn't need to export this, just a helper for DRY
function useBearContext() {
  const value = useContext(BearStoreContext);
  if (!value) {
    throw new Error('Must set up the BearStoreProvider first');
  }
  return value;
}

export function useBearState(): BearState;
export function useBearState<T>(selector: (state: BearState) => T): T;
export function useBearState<T>(selector?: (state: BearState) => T) {
  const {store} = useBearContext();
  return useStore(store, selector!);
}

// No need for a selector because this is stable across renders
export function useBearActions() {
  const {actions} = useBearContext();
  return actions;
}

Provider setup could be done as:

import React from 'react';
import {BearStoreProvider, createBearStore} from './bear-store.ts';

const bearStore = createBearStore({persist: true});

export const AppProvider = ({children}) => {
  return <BearStoreProvider value={bearStore}>{children}</BearStoreProvider>;
};

And then the store can be used in components:

import React from 'react';
import {useBearActions, useBearState} from './bear-store.ts';

function BearCounter() {
  const bears = useBearState(state => state.bears);
  return <h1>{bears} around here...</h1>;
}

function Controls() {
  const {increasePopulation} = useBearActions();
  return <button onClick={increasePopulation}>one up</button>;
}

For testing we want a different instance of the store for each test, which we can do with something like:

export const TestProvider = ({children}) => {
  const [bearStore] = React.useState(() => createBearStore({persist: false}));
  return <BearStoreProvider value={bearStore}>{children}</BearStoreProvider>;
};

Then when running tests:

describe('test 1', () => {
  render(MyComponent, { wrapper: TestProvider })
  // test something
})

describe('test 2', () => {
  render(MyComponent, { wrapper: TestProvier })
  // test something else
})

This way each test is run with a new instance of the (non-persisted) store.

If we want to expose the actual store, in order to read or mutate its value during tests, then we could do:

export function createTestContext() {
  const bearStore = createBearStore({persist: false})
  const Wrapper = ({ children }) => <BearStoreProvider value={bearStore}>{children}</BearStoreProvider>;
  return { bearStore, Wrapper }
}

Then in tests:

describe('test 1', () => {
  const { Wrapper, bearStore } = createTestContext()
  bearStore.store.setState({ bears: 100 })
  render(MyComponent, { wrapper: Wrapper })
  // test something
})

However we should probably avoid interacting with the actual store directly in tests, because it leaks an implementation detail into the tests (if we change the store structure, we want tests to still pass if they still do the same thing), so I do not think we should do this. If all we want to do is set an initial state for the store, and it really is too complicated to create that state in the test via userEvents, then we could expose an "initialState" argument for createBearStore, however I think it's much better if we create the state that we need in each test. Running tests in JSDOM should be fast, and we can create test helpers that can create the initial states we might need (like selecting a project Id). Doing things this way will result in more robust tests that allow us to change the implementation without breaking tests, which is important because we want to test the experience of the user using the app.

We might want to end up creating our own persist middleware to wrap zustand one with some defaults, however it might be better to keep things explicit and just use the persist from zustand/middleware each time, and if we need a custom storage, abstract just that instead e.g.

export function createElectronStorage() {
  return createJsonStorage(() => electronStore.getStore())
}