How to override your Zustand state.

In this article, we’ll explore how Zustand enables state replacement with the setState function, focusing on an interesting feature: overriding the state without merging. We'll explain how to use it with a simple test case.

Understanding State Merging vs State Overriding

Before looking at the code, it’s important to understand the difference between merging and overriding state.

• Merging State: By default, when you update state in Zustand (or React), it merges the new state with the existing state. Only the updated fields are changed, while the rest of the state remains intact.

• Overriding State: In contrast, overriding the state replaces the entire state with a new object, discarding any fields that were present in the previous state but are not part of the new state.

When to Override State

There are situations where replacing the entire state is necessary, for instance:

• Resetting state after a form submission.

• Replacing state entirely when loading a different dataset.

Test Case for Overriding Zustand State

Here’s a simple test case from Zustand’s codebase that demonstrates how to override the state using the setState function:

it('can set the store without merging', () => {
  const { setState, getState } = create<{ a: number } | { b: number }>(
    (_set) => ({
      a: 1,
    }),
  )

  // Should override the state instead of merging.
  setState({ b: 2 }, true)
  expect(getState()).toEqual({ b: 2 })
})

Breaking Down the Test

1. Creating the Store First, a Zustand store is created with an initial state:

const { setState, getState } = create<{ a: number } | { b: number }>(
    (_set) => ({
      a: 1,
    }),
  )

The initial state is an object with a single property a set to 1. The create function returns two essential methods:

• setState: This method is used to update the store’s state.

• getState: This method retrieves the current state of the store.

2. Overriding the State The key operation here is to replace the current state, { a: 1 }, with a completely new state:

setState({ b: 2 }, true)

• By passing the true flag as the second argument, Zustand knows that it should override the state, rather than merge the new object { b: 2 } with the existing one. This completely replaces the previous state, so { a: 1 } is removed, and the new state becomes { b: 2 }.

3. Validating the State Replacement Finally, the test ensures that the state has been fully replaced:

expect(getState()).toEqual({ b: 2 })

The expectation is that after calling setState({ b: 2 }, true), the store’s state will only contain { b: 2 } and no longer contain { a: 1 }.

How Zustand Overrides State

The below image is from Zustand source code that sets state.

As you can see, if replace flag exists, nextState become the new state.

(replace ?? (typeof nextState !== 'object' || nextState === null))

This is a clever way, if replace is false, which it is by default, (typeof nextState !== ‘object’ || nextState === null)) condition is checked.

Conclusion

Zustand offers a straightforward way to manage state in React, and the ability to override the state completely using setState with the override flag provides additional flexibility. Whether you want to reset a form, load new data, or clear out old values, this feature allows you to fully control how state is managed and updated in your application.

About us:

At Think Throo, we are on a mission to teach the advanced codebase architectural concepts used in open-source projects.

10x your coding skills by practising advanced architectural concepts in Next.js/React, learn the best practices and build production-grade projects.

We are open source — https://github.com/thinkthroo/thinkthroo (Do give us a star!)

Up skill your team with our advanced courses based on codebase architecture. Reach out to us at [e-mail] to learn more!

References:

1. https://github.com/pmndrs/zustand/blob/v4.5.5/tests/basic.test.tsx#L471C13-L471C50

2. https://github.com/pmndrs/zustand/tree/v4.5.5?tab=readme-ov-file#overwriting-state