How to mock createWithEqualityFn with Vitest?

[goliney]

The example covers the mocking of import { create } from 'zustand';

In my project I use import { createWithEqualityFn } from "zustand/traditional"; for this purpose.

My store looks like this:

import { createWithEqualityFn } from "zustand/traditional";

// ...

export const useStore = createWithEqualityFn<AppStore>()(
  immer(
    devtools(
      (...args) => {
        return {
          someSlice: createSomeSlice(...args),
          anotherSlice: createAnotherSlice(...args),
        };
      },
      { name: "My Store" }
    )
  )
);

How would you mock this so the store would be reset after each test run?

[goliney]

I came up with this:

// __mocks__/zustand/traditional.ts

import { act } from "@testing-library/react";
import type { StateCreator } from "zustand";
import type * as zustandTraditional from "zustand/traditional";

const { createWithEqualityFn: actualCreateWithEqualityFn } =
  await vi.importActual<typeof zustandTraditional>("zustand/traditional");

// a variable to hold reset functions for all stores declared in the app
export const storeResetFns = new Set<() => void>();

const createWithEqualityFnUncurried = <T>(stateCreator: StateCreator<T>) => {
  const store = actualCreateWithEqualityFn(stateCreator);
  const initialState = store.getInitialState();
  storeResetFns.add(() => {
    store.setState(initialState, true);
  });
  return store;
};

export const createWithEqualityFn = (<T>(stateCreator: StateCreator<T>) => {
  return typeof stateCreator === "function"
    ? createWithEqualityFnUncurried(stateCreator)
    : createWithEqualityFnUncurried;
}) as typeof zustandTraditional.createWithEqualityFn;

// reset all stores after each test run
afterEach(() => {
  act(() => {
    storeResetFns.forEach((resetFn) => {
      resetFn();
    });
  });
});

[dbritto-dev]

@goliney yeah, that works perfectly