Optimal Practices for Structuring Zustand Store

After spending some time playing around and testing Zustand, I’ve developed what I consider to be the best practice for structuring the use of Zustand store at the moment.

import create from 'zustand'
import { produce } from 'immer'

export type BearStateType = {
    bears: number
    name: string
}

export type BearFuncType = {
    increasePopulation: () => void
    removeAllBears: () => void
    resetBearStore: () => void
}

export type BearStoreState = BearStateType & BearFuncType

const defaultState: BearStateType = {
    bears: 9,
    name: 'NY Bears Club',
}

export const useBearStore = create((set) => {
    return {
        ...defaultState,
        increasePopulation: () =>
            set(
                produce((state) => {
                    state.bears = state.bears + 1
                })
            ),
        removeAllBears: () =>
            set(
                produce((state) => {
                    state.bears = 0
                })
            ),
        resetBearStore: () => {
            produce((state) => {
                state = defaultState
            })
        },
    }
})

export const { 
    increasePopulation, 
    removeAllBears, 
    resetBearStore 
} = useBearStore.getState() as BearFuncType

Below is the TypeScript version. I’ve successfully defined all types except for the one assigned as “any”. Although I am quite sure that the type is actually WritableDraft<T>, which is only locally accessible inside Immer.js. I’m puzzled why it’s not exported for other users to utilize, given its importance for ensuring the correctness of the code’s typings when using immer.js.

import create from 'zustand'
import { produce, Draft } from 'immer'

export type WebsiteStoreStateType = {
    hasWebsiteDataTouched: boolean
    visiters: number
}

export type WebsiteStoreFuncType = {
    setWebsiteDataTouchedState: (isTouched: boolean) => void
    setVisiters: (count: number) => void
}

export type WebsiteStoreType = WebsiteStoreStateType & WebsiteStoreFuncType

const defaultState: WebsiteStoreStateType = {
    hasWebsiteDataTouched: false,
    visiters: 0,
}

// Immerjs doesn't export the type WritableDraft<T>, 
//hence, have no choice but set the type of the 'set' argument as 'any'.
export const useWebsiteStore = create<WebsiteStoreType>((set: (arg: any) => WebsiteStoreType) => {
    return {
        ...defaultState,
        setWebsiteDataTouchedState: (isTouched: boolean) => {
            const newState = produce((draftState: Draft<WebsiteStoreType>) => {
                draftState.hasWebsiteDataTouched = isTouched
            })
            set(newState)
        },
        setVisiters: (count: number) => {
            const newState = produce((draftState: Draft<WebsiteStoreType>) => {
                draftState.count = count
            })
            set(newState)
        },
    }
})

export const { 
    setWebsiteDataTouchedState, 
    setVisiters 
} = useWebsiteStore.getState() as WebsiteStoreFuncType

Opt for Immer Over Zukeeper

While Zukeeper works perfectly with Zustand, effectively handling immutable updates for nested states, it introduces a critical issue during component testing with React Testing Library. This issue leads to the Zustand store becoming undefined or null in the testing environment, resulting in multiple failures all over the places in the tests. And there is no easy fixes that I know of. On the other hand, using Immer resolves this issue while maintaining support for immutable updates in nested states.

Avoid Retrieving Functions From the Store Hook

Retrieving functions directly from the store hook can be as simple as a one-liner at the top of components.

const YourApp = (props) => {
    const { increasePopulation } = useBearStore(state => state.increasePopulation)
    ...
    
    render( ...

This approach enables you to utilize these functions within components for events such as onClick or onKeyUp … and etc.. However, it results in unnecessary re-renders of the component. This is because the hook continuously monitors changes to the state property, and since a function is an object in Javascript, each time the component is executed, the function’s address changes. As a result, the component is forced to re-render every time, even when there are no distinguishable changes in the state visible to the human eye.

Instead, we should obtain the functions from the store’s state using getState(). This approach allows you to utilize these functions within components for events without introducing an observer to constantly monitor changes in the function’s address. This simple fix easily eliminates unnecessary re-renders.

To take this one step further, we can enhance the code by directly exporting the functions from the Zustand store, as they remain unchanged throughout the entire runtime of the app. This is demonstrated in the last few lines of the sample code provided at the beginning of this post.

Adding a Default State

Having a default state helps in resetting the state as necessary. With a refreshed state (only if needed) each time a component is unmounted and remounted, it simplifies the data flow within the app.

Separating Data and Function Types

This tiny improvement keeps the store cleaner and prevents accidentally adding data to the export list, which could cause minor issues for components.

Conclusion

Zustand is a lightweight state management framework that offers great flexibility in code writing. However, employing a well-structured pattern enhances its effectiveness and makes it more suitable for medium to large-sized applications.