Mocking zustand store for jest test - reactjs

Question

hej, maybe someone can help me with this? I try to mock a zustand store with custom values and actions. I'd like to do the mocking within my test file. But I fail miserably.
This is a simplified example of my react app.
component.jsx
import React from "react";
import useStore from "#app/component_store";
export default function Component() {
const foo = useStore((state) => state.foo);
const setFoo = useStore((state) => state.foo_set);
return (
<>
<div data-testid="foo">{foo}</div>
<button data-testid="btn-foo" onClick={() => setFoo("new foo")}>
set foo
</button>
</>
);
}
component_store.jsx
import create from "zustand";
const useStore = create((set) => ({
foo: "foo",
foo_set: (value) => set(() => ({ foo: value })),
}));
export default useStore;
component.test.jsx
/* global afterEach, expect, jest, test */
import React from "react";
import { cleanup, fireEvent, render } from "#testing-library/react";
import "#testing-library/jest-dom";
import mockCreate from "zustand";
import Component from "#app/component";
jest.mock("#app/component_store", () => (args) => {
const useStore = mockCreate((set, get) => ({
foo: "mocked foo",
foo_set: (value) => set(() => ({ foo: "mocked set foo" })),
}));
return useStore(args);
});
afterEach(cleanup);
test("override store values and actions", () => {
const { container, getByTestId } = render(<Component />);
const foo = getByTestId("foo");
const btnFoo = getByTestId("btn-foo");
expect(foo.textContent).toBe("mocked foo");
fireEvent.click(btnFoo);
expect(foo.textContent).toBe("mocked set foo");
});
This fails. expect(foo.textContent).toBe("mocked set foo"); will still have "mocked foo" as a value, although the mocked foo_set gets called (e.g. adding a console log does get logged). But again, the value of foo will not be updated.
However, if I move the store mocking part into its own file and import it into my test, everything will work as expected.
component_store_mocked.jsx
note: this is the same as in component.test.jsx
import create from "zustand";
const useStore = create((set, get) => ({
foo: "mocked foo",
foo_set: (value) => set(() => ({ foo: "mocked set foo" })),
}));
export default useStore;
updated component.test.jsx
/* global afterEach, expect, jest, test */
import React from "react";
import { cleanup, fireEvent, render } from "#testing-library/react";
import "#testing-library/jest-dom";
import mockCreate from "zustand";
import Component from "#app/component";
// this part here
import mockComponentStore from "#app/component_store_mocked";
jest.mock("#app/component_store", () => (args) => {
const useStore = mockComponentStore;
return useStore(args);
});
afterEach(cleanup);
test("override store values and actions", () => {
const { container, getByTestId } = render(<Component />);
const foo = getByTestId("foo");
const btnFoo = getByTestId("btn-foo");
expect(foo.textContent).toBe("mocked foo");
fireEvent.click(btnFoo);
expect(foo.textContent).toBe("mocked set foo");
});
I would love to do the mocking within the test file and not pollute my system with dozens of mocked stores for my test ... Again, this is just a simplified setup of a much larger app.