diff --git a/tests/__mocks__/zustand/index.ts b/tests/__mocks__/zustand/index.ts
index fbb962d..d085b8b 100644
--- a/tests/__mocks__/zustand/index.ts
+++ b/tests/__mocks__/zustand/index.ts
@@ -1,22 +1,54 @@
-import { type StateCreator, create as actualCreate } from 'zustand';
-import { act } from 'react-dom/test-utils';
-import { afterEach } from 'vitest';
+// https://github.com/pmndrs/zustand/blob/main/docs/guides/testing.md
+import * as zustand from 'zustand';
+import { act } from '@testing-library/react';
+
+const { create: actualCreate, createStore: actualCreateStore } =
+  await vi.importActual<typeof zustand>('zustand');
 
 // a variable to hold reset functions for all stores declared in the app
-const storeResetFns = new Set<() => void>();
+export const storeResetFns = new Set<() => void>();
+
+const createUncurried = <T>(stateCreator: zustand.StateCreator<T>) => {
+  const store = actualCreate(stateCreator);
+  const initialState = store.getInitialState();
+  storeResetFns.add(() => {
+    store.setState(initialState, true);
+  });
+  return store;
+};
 
 // when creating a store, we get its initial state, create a reset function and add it in the set
-const createStore = (createState: StateCreator<unknown>) => {
-  if (!createState) return createStore;
-  const store = actualCreate(createState);
-  const initialState = store.getState();
-  storeResetFns.add(() => store.setState(initialState, true));
+export const create = (<T>(stateCreator: zustand.StateCreator<T>) => {
+  // to support curried version of create
+  return typeof stateCreator === 'function'
+    ? createUncurried(stateCreator)
+    : createUncurried;
+}) as typeof zustand.create;
+
+const createStoreUncurried = <T>(stateCreator: zustand.StateCreator<T>) => {
+  const store = actualCreateStore(stateCreator);
+  const initialState = store.getInitialState();
+  storeResetFns.add(() => {
+    store.setState(initialState, true);
+  });
   return store;
 };
 
-// Reset all stores after each test run
+// when creating a store, we get its initial state, create a reset function and add it in the set
+export const createStore = (<T>(stateCreator: zustand.StateCreator<T>) => {
+  // to support curried version of createStore
+  return typeof stateCreator === 'function'
+    ? createStoreUncurried(stateCreator)
+    : createStoreUncurried;
+}) as typeof zustand.createStore;
+
+export const { useStore } = zustand;
+
+// reset all stores after each test run
 afterEach(() => {
-  act(() => storeResetFns.forEach((resetFn) => resetFn()));
+  act(() => {
+    storeResetFns.forEach((resetFn) => {
+      resetFn();
+    });
+  });
 });
-
-export { createStore };
diff --git a/tests/__tests__/createVanillaTemporal.test.ts b/tests/__tests__/createVanillaTemporal.test.ts
index c70ef5e..faf000e 100644
--- a/tests/__tests__/createVanillaTemporal.test.ts
+++ b/tests/__tests__/createVanillaTemporal.test.ts
@@ -1,8 +1,7 @@
-import { describe, it, expect, vi } from 'vitest';
-vi.mock('zustand');
+import { describe, it, expect } from 'vitest';
 import { temporalStateCreator } from '../../src/temporal';
 import { createStore } from 'zustand';
-import { act } from 'react-dom/test-utils';
+import { act } from '@testing-library/react';
 
 interface MyState {
   count: number;
diff --git a/tests/__tests__/options.test.ts b/tests/__tests__/options.test.ts
index b2b318d..e294590 100644
--- a/tests/__tests__/options.test.ts
+++ b/tests/__tests__/options.test.ts
@@ -1,9 +1,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-
-vi.mock('zustand');
 import { temporal } from '../../src/index';
 import { createStore, type StoreApi } from 'zustand';
-import { act } from 'react-dom/test-utils';
+import { act } from '@testing-library/react';
 import { shallow } from 'zustand/shallow';
 import type {
   _TemporalState,
diff --git a/tests/__tests__/react.test.tsx b/tests/__tests__/react.test.tsx
index 9ff4833..61178f1 100644
--- a/tests/__tests__/react.test.tsx
+++ b/tests/__tests__/react.test.tsx
@@ -1,12 +1,9 @@
 import { describe, it, expect } from 'vitest';
-import React from 'react';
 import { fireEvent, render } from '@testing-library/react';
 
 describe('React Re-renders when state changes', () => {
   it('it', () => {
-    const { queryByLabelText, getByLabelText, queryByText, getByText } = render(
-      <Reactive />,
-    );
+    const { queryByText, getByText } = render(<Reactive />);
 
     expect(queryByText(/bears: 0/i)).toBeTruthy();
     expect(queryByText(/increment/i)).toBeTruthy();
@@ -77,9 +74,8 @@ const useTemporalStore = <
   U,
 >(
   selector: (state: ExtractState<S>) => U,
-  equality?: (a: U, b: U) => boolean,
 ): U => {
-  const state = useStore(useMyStore.temporal as any, selector, equality);
+  const state = useStore(useMyStore.temporal as any, selector);
   return state;
 };
 
@@ -97,10 +93,8 @@ const HistoryBar = () => {
 };
 
 const UndoBar = () => {
-  const { undo, redo } = useTemporalStore((state) => ({
-    undo: state.undo,
-    redo: state.redo,
-  }));
+  const undo = useTemporalStore((state) => state.undo);
+  const redo = useTemporalStore((state) => state.redo);
   return (
     <div>
       <button onClick={() => undo()}>undo</button>
diff --git a/tests/__tests__/zundo.test.ts b/tests/__tests__/zundo.test.ts
index ac923d5..45fcc86 100644
--- a/tests/__tests__/zundo.test.ts
+++ b/tests/__tests__/zundo.test.ts
@@ -1,8 +1,7 @@
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-vi.mock('zustand');
 import { temporal } from '../../src/index';
 import { createStore, type StoreApi } from 'zustand';
-import { act } from 'react-dom/test-utils';
+import { act } from '@testing-library/react';
 import type { TemporalState, Write } from '../../src/types';
 
 interface MyState {
diff --git a/tests/tsconfig.json b/tests/tsconfig.json
index 861c823..9df0421 100644
--- a/tests/tsconfig.json
+++ b/tests/tsconfig.json
@@ -13,5 +13,6 @@
     "resolveJsonModule": true,
     "esModuleInterop": true,
     "moduleResolution": "Bundler",
+    "types": ["vitest/globals"]
   }
 }
diff --git a/tests/vitest.config.ts b/tests/vitest.config.ts
index bb921f4..2a4eef2 100644
--- a/tests/vitest.config.ts
+++ b/tests/vitest.config.ts
@@ -7,7 +7,6 @@ export default defineConfig({
     environment: 'happy-dom',
     dir: '.',
     setupFiles: ['./vitest.setup.ts', 'vitest-localstorage-mock'],
-    mockReset: false,
     coverage: {
       include: ['**/src/**/*.{ts,tsx}'],
       allowExternal: true,
diff --git a/tests/vitest.setup.ts b/tests/vitest.setup.ts
index 0d9bbfe..2e90f04 100644
--- a/tests/vitest.setup.ts
+++ b/tests/vitest.setup.ts
@@ -1,8 +1,4 @@
-import { afterEach } from 'vitest';
-import { cleanup } from '@testing-library/react';
-import '@testing-library/jest-dom/vitest';
+// https://github.com/pmndrs/zustand/blob/main/docs/guides/testing.md
+import '@testing-library/jest-dom';
 
-// runs a cleanup after each test case (e.g. clearing happy-dom)
-afterEach(() => {
-  cleanup();
-});
+vi.mock('zustand'); // to make it work like Jest (auto-mocking)