MobX-State-Tree Error: Mocking action returns vi.spy error

I have a React/Typescript project that I'm using MobX-State-Tree for state management. In one of my tests (using Vitest + React Testing Library), I mocked a function defined in one of the stores in the state tree, so I can test for its call (it should get called in a useEffect on mount). But the test is not working correctly.

Sample store:

const SampleStore = types.model({
   ...
}).actions(self => ({
   sample_action: flow(function* (param: string) {
     //sample async action call with yield
   })
})
...
export const sample_store = SampleStore.create()

Component:

// all imports

export const SampleComponentPage = () => {
  ....

  useEffect(() => {
    // the action from the store gets called here
    sample_store.sample_action(arg)
  }, [])

  return (
    ...
  )
})

Test file:

...

describe('Sample component page', () => {
  it('sample test case', async () => {
    const mock_action = vi.fn()

    const initialStore = SampleStore.create({
       sample_action: mock_action
    })

    render(
      <StoreContextProvider value={initialStore}>
        <SampleComponentPage />
      </StoreContextProvider>
    )

    ...
    waitFor(() => expect(mock_action).toHaveBeenCalled())
    ...
  })
})

The waitFor statement is returning a false positive test - the test passes regardless e.g with toHaveBeenCalledTimes(10) even though it should only have been called once.

Without the waitFor i.e calling expect(mock_action).toHaveBeenCalled(), it returns this error:

AssertionError: expected "spy" to be called once, but got 0 times

I've tried using vi.spyOn(initialStore, 'sample_action') for pointing the store action to the mock function, but that returns a MST-related error:

Error: [mobx-state-tree] Cannot modify 'AnonymousModel@/...', the object is protected and can only be modified by using an action.
 ❯ fail node_modules/mobx-state-tree/dist/mobx-state-tree.js:3901:12

What's the best fix for this please? I want to test that the store action gets called.

I encountered a similar issue when migrating our test suite from jest to vitest. The model is protected by MobX-State-Tree by default and can't be modified directly just like the error says. However you can turn this off in your test by calling unprotect(initialStore) before vi.spyOn.

import { unprotect } from "mobx-state-tree";

it('sample test case', async () => {
  const initialStore = SampleStore.create(...);
  unprotect(initialStore);
  const spy = vi.spyOn(initialStore, 'submit_answer')
  expect(spy).toHaveBeenCalled();
});

Check out MobX-State-Tree Unprotect API documentation.