mocked mutateAsync should call actual onSuccess callback while mocking useMutation using Vitest

(Next.Js, react-query, Vitest)

I have a modal where the user has a few steps to complete. At the last step, code sends a request to the backend. This request is sent using useMutation().mutateAsync and in useMutation().onSuccess I handle showing the success screen.

I aim to mock my useUpload hook to test with what parameters it has been called.

const mockedUpload = vi.fn();
vi.mock('hooks/useUpload.tsx', () => ({
  useUpload: () => ({
    mutateAsync: mockedUpload,
  }),
}));

// ...


expect(mockedMutate).toHaveBeenCalledWith({ test: '123' })

It works, but unfortunately, it breaks showing the success screen since it is set from onSuccess which in this mocked version is never called. Ideally mocked mutateAsync should call actual onSuccess but I didn't manage the mocking code.

Any ideas?

The useUpload hook looks like this:

export const useUpload = ({ onSuccess }: UseUploadProps = {}) => {
  const mutation = useMutation({
    mutationFn: async (props: UploadProps) => uploadFile(props),
    onSuccess() {
      onSuccess?.();
    },
  });

  return { ...mutation, isLoading: mutation.isPending };
};

I found the other way, which takes passed onSuccess and calls it directly from mutateAsync instead of mocking whole useMutation so that it's mutateAsync should call actual onSuccess which should call passed onSuccess. Finally mock looks like this:

vi.mock('hooks/useUpload.tsx', () => ({
  useUpload: ({ onSuccess }: { onSuccess: () => void }) => ({
    mutateAsync: vi.fn().mockImplementation(() => {
      onSuccess();
    }),
    onSuccess: vi.fn(),
  }),
}));

and to separate it as a variable:

// Define a variable to store the onSuccess callback
let onSuccessCallback = vi.fn();

const mockMutateAsync = vi.fn().mockImplementation(() => {
  onSuccessCallback?.();
});

const mockOnSuccess = vi.fn();

// Mock the useUpload hook
vi.mock('hooks/useUpload.tsx', () => ({
  useUpload({ onSuccess }: { onSuccess: typeof mockOnSuccess }) {
    // Store the provided onSuccess callback in a variable
    onSuccessCallback = onSuccess;
    return {
      mutateAsync: mockMutateAsync,
      onSuccess: mockOnSuccess,
    };
  },
}));

and usage like:

expect(mockMutateAsync).toHaveBeenCalledWith({
  axiosConfig: {
    headers: {
      'Content-Type': 'multipart/form-data',
    },
  },
  formData: expect.any(FormData) as FormData,
  // ... other params to test
});

Are there any better ways of fixing the mentioned issue? If so, I would love to hear and learn.