From 5377b04911627da987cf2c388c6695452723cb19 Mon Sep 17 00:00:00 2001
From: Gildas Garcia <1122076+djhi@users.noreply.github.com>
Date: Thu, 23 May 2024 17:50:21 +0200
Subject: [PATCH 1/3] Refactor middlewares to handle optimistic cases

---
 .../controller/create/useCreateController.ts  |   4 +-
 .../src/controller/edit/useEditController.ts  |   4 +-
 .../saveContext/useMutationMiddlewares.ts     |  33 +-
 .../src/dataProvider/useCreate.spec.tsx       |  49 +++
 .../src/dataProvider/useCreate.stories.tsx    | 307 ++++++++++++++++++
 .../ra-core/src/dataProvider/useCreate.ts     |  24 +-
 .../useUpdate.optimistic.stories.tsx          |   4 +-
 .../useUpdate.pessimistic.stories.tsx         |   4 +-
 .../ra-core/src/dataProvider/useUpdate.ts     |  35 +-
 .../useUpdate.undoable.stories.tsx            |   4 +-
 10 files changed, 411 insertions(+), 57 deletions(-)
 create mode 100644 packages/ra-core/src/dataProvider/useCreate.stories.tsx

diff --git a/packages/ra-core/src/controller/create/useCreateController.ts b/packages/ra-core/src/controller/create/useCreateController.ts
index 2a5820544b3..f18bba241d5 100644
--- a/packages/ra-core/src/controller/create/useCreateController.ts
+++ b/packages/ra-core/src/controller/create/useCreateController.ts
@@ -79,7 +79,7 @@ export const useCreateController = <
         mutationOptions;
     const {
         registerMutationMiddleware,
-        mutateWithMiddlewares,
+        getMutateWithMiddlewares,
         unregisterMutationMiddleware,
     } = useMutationMiddlewares();
 
@@ -134,7 +134,7 @@ export const useCreateController = <
         },
         ...otherMutationOptions,
         returnPromise: true,
-        mutateWithMiddlewares,
+        getMutateWithMiddlewares,
     });
 
     const save = useCallback(
diff --git a/packages/ra-core/src/controller/edit/useEditController.ts b/packages/ra-core/src/controller/edit/useEditController.ts
index 7ea7e87ab4a..96c244ef3e9 100644
--- a/packages/ra-core/src/controller/edit/useEditController.ts
+++ b/packages/ra-core/src/controller/edit/useEditController.ts
@@ -91,7 +91,7 @@ export const useEditController = <
     } = mutationOptions;
     const {
         registerMutationMiddleware,
-        mutateWithMiddlewares,
+        getMutateWithMiddlewares,
         unregisterMutationMiddleware,
     } = useMutationMiddlewares();
     const {
@@ -192,7 +192,7 @@ export const useEditController = <
             ...otherMutationOptions,
             mutationMode,
             returnPromise: mutationMode === 'pessimistic',
-            mutateWithMiddlewares,
+            getMutateWithMiddlewares,
         }
     );
 
diff --git a/packages/ra-core/src/controller/saveContext/useMutationMiddlewares.ts b/packages/ra-core/src/controller/saveContext/useMutationMiddlewares.ts
index 0c6538c3859..5b744f0c427 100644
--- a/packages/ra-core/src/controller/saveContext/useMutationMiddlewares.ts
+++ b/packages/ra-core/src/controller/saveContext/useMutationMiddlewares.ts
@@ -51,12 +51,11 @@ export const useMutationMiddlewares = <
         []
     );
 
-    const mutateWithMiddlewares = useCallback(
-        (
-            fn: MutateFunc,
-            ...args: Parameters<MutateFunc>
-        ): ReturnType<MutateFunc> => {
-            let index = callbacks.current.length - 1;
+    const getMutateWithMiddlewares = useCallback((fn: MutateFunc) => {
+        // Stores the current callbacks in a closure to avoid losing them if the calling component is unmounted
+        const currentCallbacks = [...callbacks.current];
+        return (...args: Parameters<MutateFunc>): ReturnType<MutateFunc> => {
+            let index = currentCallbacks.length - 1;
 
             // Called by middlewares to call the next middleware function
             // Should take the same arguments as the original mutation function
@@ -67,32 +66,31 @@ export const useMutationMiddlewares = <
 
                 // If there are no more middlewares, we call the original mutation function
                 if (index >= 0) {
-                    return callbacks.current[index](...newArgs, next);
+                    return currentCallbacks[index](...newArgs, next);
                 } else {
                     return fn(...newArgs);
                 }
             };
 
-            if (callbacks.current.length > 0) {
+            if (currentCallbacks.length > 0) {
                 // Call the first middleware with the same args as the original mutation function
                 // with an additional next function
-                return callbacks.current[index](...args, next);
+                return currentCallbacks[index](...args, next);
             }
 
             return fn(...args);
-        },
-        []
-    );
+        };
+    }, []);
 
     const functions = useMemo<UseMutationMiddlewaresResult<MutateFunc>>(
         () => ({
             registerMutationMiddleware,
-            mutateWithMiddlewares,
+            getMutateWithMiddlewares,
             unregisterMutationMiddleware,
         }),
         [
             registerMutationMiddleware,
-            mutateWithMiddlewares,
+            getMutateWithMiddlewares,
             unregisterMutationMiddleware,
         ]
     );
@@ -104,10 +102,9 @@ export interface UseMutationMiddlewaresResult<
     MutateFunc extends (...args: any[]) => any = (...args: any[]) => any,
 > {
     registerMutationMiddleware: (callback: Middleware<MutateFunc>) => void;
-    mutateWithMiddlewares: (
-        mutate: MutateFunc,
-        ...args: Parameters<MutateFunc>
-    ) => ReturnType<MutateFunc>;
+    getMutateWithMiddlewares: (
+        mutate: MutateFunc
+    ) => (...args: Parameters<MutateFunc>) => ReturnType<MutateFunc>;
     unregisterMutationMiddleware: (callback: Middleware<MutateFunc>) => void;
 }
 
diff --git a/packages/ra-core/src/dataProvider/useCreate.spec.tsx b/packages/ra-core/src/dataProvider/useCreate.spec.tsx
index e64b5d6df33..0efcaf5ad7a 100644
--- a/packages/ra-core/src/dataProvider/useCreate.spec.tsx
+++ b/packages/ra-core/src/dataProvider/useCreate.spec.tsx
@@ -7,6 +7,10 @@ import { testDataProvider } from './testDataProvider';
 import { useCreate } from './useCreate';
 import { useGetList } from './useGetList';
 import { CoreAdminContext } from '../core';
+import {
+    WithMiddlewaresError,
+    WithMiddlewaresSuccess,
+} from './useCreate.stories';
 
 describe('useCreate', () => {
     it('returns a callback that can be used with create arguments', async () => {
@@ -181,4 +185,49 @@ describe('useCreate', () => {
             expect(screen.queryByText('ghi')).not.toBeNull();
         });
     });
+
+    describe('middlewares', () => {
+        it('it accepts middlewares and displays result and success side effects when dataProvider promise resolves', async () => {
+            render(<WithMiddlewaresSuccess timeout={10} />);
+            screen.getByText('Create post').click();
+            await waitFor(() => {
+                expect(screen.queryByText('success')).toBeNull();
+                expect(
+                    screen.queryByText('Hello World from middleware')
+                ).toBeNull();
+                expect(screen.queryByText('mutating')).not.toBeNull();
+            });
+            await waitFor(() => {
+                expect(screen.queryByText('success')).not.toBeNull();
+                expect(
+                    screen.queryByText('Hello World from middleware')
+                ).not.toBeNull();
+                expect(screen.queryByText('mutating')).toBeNull();
+            });
+        });
+
+        it('it accepts middlewares and displays error and error side effects when dataProvider promise rejects', async () => {
+            jest.spyOn(console, 'error').mockImplementation(() => {});
+            render(<WithMiddlewaresError timeout={10} />);
+            screen.getByText('Create post').click();
+            await waitFor(() => {
+                expect(screen.queryByText('success')).toBeNull();
+                expect(screen.queryByText('something went wrong')).toBeNull();
+                expect(
+                    screen.queryByText('Hello World from middleware')
+                ).toBeNull();
+                expect(screen.queryByText('mutating')).not.toBeNull();
+            });
+            await waitFor(() => {
+                expect(screen.queryByText('success')).toBeNull();
+                expect(
+                    screen.queryByText('something went wrong')
+                ).not.toBeNull();
+                expect(
+                    screen.queryByText('Hello World from middleware')
+                ).toBeNull();
+                expect(screen.queryByText('mutating')).toBeNull();
+            });
+        });
+    });
 });
diff --git a/packages/ra-core/src/dataProvider/useCreate.stories.tsx b/packages/ra-core/src/dataProvider/useCreate.stories.tsx
new file mode 100644
index 00000000000..3d97dc26729
--- /dev/null
+++ b/packages/ra-core/src/dataProvider/useCreate.stories.tsx
@@ -0,0 +1,307 @@
+import * as React from 'react';
+import { useState } from 'react';
+import { QueryClient, useIsMutating } from '@tanstack/react-query';
+
+import { CoreAdminContext } from '../core';
+import { useCreate } from './useCreate';
+import { useGetOne } from './useGetOne';
+
+export default { title: 'ra-core/dataProvider/useCreate' };
+
+export const SuccessCase = ({ timeout = 1000 }) => {
+    const posts: { id: number; title: string; author: string }[] = [];
+    const dataProvider = {
+        getOne: (resource, params) => {
+            return Promise.resolve({
+                data: posts.find(p => p.id === params.id),
+            });
+        },
+        create: (resource, params) => {
+            return new Promise(resolve => {
+                setTimeout(() => {
+                    const post = { id: posts.length + 1, ...params.data };
+                    posts.push(post);
+                    resolve({ data: post });
+                }, timeout);
+            });
+        },
+    } as any;
+    return (
+        <CoreAdminContext
+            queryClient={new QueryClient()}
+            dataProvider={dataProvider}
+        >
+            <SuccessCore />
+        </CoreAdminContext>
+    );
+};
+
+const SuccessCore = () => {
+    const isMutating = useIsMutating();
+    const [success, setSuccess] = useState<string>();
+    const { data, refetch } = useGetOne(
+        'posts',
+        { id: 1 },
+        { enabled: success === 'success' }
+    );
+    const [create, { isPending }] = useCreate();
+    const handleClick = () => {
+        create(
+            'posts',
+            {
+                data: { title: 'Hello World' },
+            },
+            {
+                onSuccess: () => setSuccess('success'),
+            }
+        );
+    };
+    return (
+        <>
+            <dl>
+                <dt>title</dt>
+                <dd>{data?.title}</dd>
+            </dl>
+            <div>
+                <button onClick={handleClick} disabled={isPending}>
+                    Create post
+                </button>
+                &nbsp;
+                <button onClick={() => refetch()}>Refetch</button>
+            </div>
+            {success && <div>{success}</div>}
+            {isMutating !== 0 && <div>mutating</div>}
+        </>
+    );
+};
+
+export const ErrorCase = ({ timeout = 1000 }) => {
+    const posts: { id: number; title: string; author: string }[] = [];
+    const dataProvider = {
+        getOne: (resource, params) => {
+            return Promise.resolve({
+                data: posts.find(p => p.id === params.id),
+            });
+        },
+        create: () => {
+            return new Promise((resolve, reject) => {
+                setTimeout(() => {
+                    reject(new Error('something went wrong'));
+                }, timeout);
+            });
+        },
+    } as any;
+    return (
+        <CoreAdminContext
+            queryClient={new QueryClient()}
+            dataProvider={dataProvider}
+        >
+            <ErrorCore />
+        </CoreAdminContext>
+    );
+};
+
+const ErrorCore = () => {
+    const isMutating = useIsMutating();
+    const [success, setSuccess] = useState<string>();
+    const [error, setError] = useState<any>();
+    const { data, refetch } = useGetOne(
+        'posts',
+        { id: 1 },
+        { enabled: success === 'success' }
+    );
+    const [create, { isPending }] = useCreate();
+    const handleClick = () => {
+        setError(undefined);
+        create(
+            'posts',
+            {
+                data: { title: 'Hello World' },
+            },
+            {
+                onSuccess: () => setSuccess('success'),
+                onError: e => setError(e),
+            }
+        );
+    };
+    return (
+        <>
+            <dl>
+                <dt>title</dt>
+                <dd>{data?.title}</dd>
+            </dl>
+            <div>
+                <button onClick={handleClick} disabled={isPending}>
+                    Create post
+                </button>
+                &nbsp;
+                <button onClick={() => refetch()}>Refetch</button>
+            </div>
+            {error && <div>{error.message}</div>}
+            {success && <div>{success}</div>}
+            {isMutating !== 0 && <div>mutating</div>}
+        </>
+    );
+};
+
+export const WithMiddlewaresSuccess = ({ timeout = 1000 }) => {
+    const posts: { id: number; title: string; author: string }[] = [];
+    const dataProvider = {
+        getOne: (resource, params) => {
+            return Promise.resolve({
+                data: posts.find(p => p.id === params.id),
+            });
+        },
+        create: (resource, params) => {
+            return new Promise(resolve => {
+                setTimeout(() => {
+                    const post = { id: posts.length + 1, ...params.data };
+                    posts.push(post);
+                    resolve({ data: post });
+                }, timeout);
+            });
+        },
+    } as any;
+    return (
+        <CoreAdminContext
+            queryClient={new QueryClient()}
+            dataProvider={dataProvider}
+        >
+            <WithMiddlewaresSuccessCore />
+        </CoreAdminContext>
+    );
+};
+
+const WithMiddlewaresSuccessCore = () => {
+    const isMutating = useIsMutating();
+    const [success, setSuccess] = useState<string>();
+    const { data, refetch } = useGetOne(
+        'posts',
+        { id: 1 },
+        { enabled: success === 'success' }
+    );
+    const [create, { isPending }] = useCreate(
+        'posts',
+        {
+            data: { title: 'Hello World' },
+        },
+        {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
+                return mutate(resource, {
+                    ...params,
+                    data: { title: `${params.data.title} from middleware` },
+                });
+            },
+        }
+    );
+    const handleClick = () => {
+        create(
+            'posts',
+            {
+                data: { title: 'Hello World' },
+            },
+            {
+                onSuccess: () => setSuccess('success'),
+            }
+        );
+    };
+    return (
+        <>
+            <dl>
+                <dt>title</dt>
+                <dd>{data?.title}</dd>
+            </dl>
+            <div>
+                <button onClick={handleClick} disabled={isPending}>
+                    Create post
+                </button>
+                &nbsp;
+                <button onClick={() => refetch()}>Refetch</button>
+            </div>
+            {success && <div>{success}</div>}
+            {isMutating !== 0 && <div>mutating</div>}
+        </>
+    );
+};
+
+export const WithMiddlewaresError = ({ timeout = 1000 }) => {
+    const posts: { id: number; title: string; author: string }[] = [];
+    const dataProvider = {
+        getOne: (resource, params) => {
+            return Promise.resolve({
+                data: posts.find(p => p.id === params.id),
+            });
+        },
+        create: () => {
+            return new Promise((resolve, reject) => {
+                setTimeout(() => {
+                    reject(new Error('something went wrong'));
+                }, timeout);
+            });
+        },
+    } as any;
+    return (
+        <CoreAdminContext
+            queryClient={new QueryClient()}
+            dataProvider={dataProvider}
+        >
+            <WithMiddlewaresErrorCore />
+        </CoreAdminContext>
+    );
+};
+
+const WithMiddlewaresErrorCore = () => {
+    const isMutating = useIsMutating();
+    const [success, setSuccess] = useState<string>();
+    const [error, setError] = useState<any>();
+    const { data, refetch } = useGetOne(
+        'posts',
+        { id: 1 },
+        { enabled: success === 'success' }
+    );
+    const [create, { isPending }] = useCreate(
+        'posts',
+        {
+            data: { title: 'Hello World' },
+        },
+        {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
+                return mutate(resource, {
+                    ...params,
+                    data: { title: `${params.data.title} from middleware` },
+                });
+            },
+        }
+    );
+    const handleClick = () => {
+        setError(undefined);
+        create(
+            'posts',
+            {
+                data: { title: 'Hello World' },
+            },
+            {
+                onSuccess: () => setSuccess('success'),
+                onError: e => setError(e),
+            }
+        );
+    };
+    return (
+        <>
+            <dl>
+                <dt>title</dt>
+                <dd>{data?.title}</dd>
+            </dl>
+            <div>
+                <button onClick={handleClick} disabled={isPending}>
+                    Create post
+                </button>
+                &nbsp;
+                <button onClick={() => refetch()}>Refetch</button>
+            </div>
+            {error && <div>{error.message}</div>}
+            {success && <div>{success}</div>}
+            {isMutating !== 0 && <div>mutating</div>}
+        </>
+    );
+};
diff --git a/packages/ra-core/src/dataProvider/useCreate.ts b/packages/ra-core/src/dataProvider/useCreate.ts
index 04c0fa888b6..271b2953cde 100644
--- a/packages/ra-core/src/dataProvider/useCreate.ts
+++ b/packages/ra-core/src/dataProvider/useCreate.ts
@@ -85,7 +85,7 @@ export const useCreate = <
     const hasCallTimeOnError = useRef(false);
     const hasCallTimeOnSuccess = useRef(false);
     const hasCallTimeOnSettled = useRef(false);
-    const { mutateWithMiddlewares, ...mutationOptions } = options;
+    const { getMutateWithMiddlewares, ...mutationOptions } = options;
     const mutation = useMutation<
         ResultRecordType,
         MutationError,
@@ -106,15 +106,14 @@ export const useCreate = <
                     'useCreate mutation requires a non-empty data object'
                 );
             }
-            if (mutateWithMiddlewares) {
-                return mutateWithMiddlewares(
-                    dataProvider.create.bind(dataProvider),
-                    callTimeResource,
-                    {
-                        data: callTimeData,
-                        meta: callTimeMeta,
-                    }
-                ).then(({ data }) => data);
+            if (getMutateWithMiddlewares) {
+                const createWithMiddlewares = getMutateWithMiddlewares(
+                    dataProvider.create.bind(dataProvider)
+                );
+                return createWithMiddlewares(callTimeResource, {
+                    data: callTimeData,
+                    meta: callTimeMeta,
+                }).then(({ data }) => data);
             }
             return dataProvider
                 .create<RecordType, ResultRecordType>(callTimeResource, {
@@ -226,11 +225,12 @@ export type UseCreateOptions<
     'mutationFn'
 > & {
     returnPromise?: boolean;
-    mutateWithMiddlewares?: <
+    getMutateWithMiddlewares?: <
         CreateFunctionType extends
             DataProvider['create'] = DataProvider['create'],
     >(
-        mutate: CreateFunctionType,
+        mutate: CreateFunctionType
+    ) => (
         ...Params: Parameters<CreateFunctionType>
     ) => ReturnType<CreateFunctionType>;
 };
diff --git a/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx b/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx
index b6559c6a1df..fa92f99626c 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx
@@ -190,7 +190,7 @@ const WithMiddlewaresSuccessCore = () => {
         },
         {
             mutationMode: 'optimistic',
-            mutateWithMiddlewares: async (mutate, resource, params) => {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
                     data: { title: `${params.data.title} from middleware` },
@@ -270,7 +270,7 @@ const WithMiddlewaresErrorCore = () => {
         },
         {
             mutationMode: 'optimistic',
-            mutateWithMiddlewares: async (mutate, resource, params) => {
+            mutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
                     data: { title: `${params.data.title} from middleware` },
diff --git a/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx b/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx
index bd2afd70ac9..a86c5c4ce6c 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx
@@ -188,7 +188,7 @@ const WithMiddlewaresSuccessCore = () => {
         },
         {
             mutationMode: 'pessimistic',
-            mutateWithMiddlewares: async (mutate, resource, params) => {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
                     data: { title: `${params.data.title} from middleware` },
@@ -269,7 +269,7 @@ const WithMiddlewaresErrorCore = () => {
         },
         {
             mutationMode: 'pessimistic',
-            mutateWithMiddlewares: async (mutate, resource, params) => {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
                     data: { title: `${params.data.title} from middleware` },
diff --git a/packages/ra-core/src/dataProvider/useUpdate.ts b/packages/ra-core/src/dataProvider/useUpdate.ts
index ff592a2e4c2..8fda6893108 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.ts
+++ b/packages/ra-core/src/dataProvider/useUpdate.ts
@@ -94,12 +94,14 @@ export const useUpdate = <RecordType extends RaRecord = any, ErrorType = Error>(
     const { id, data, meta } = params;
     const {
         mutationMode = 'pessimistic',
-        mutateWithMiddlewares,
+        getMutateWithMiddlewares,
         ...mutationOptions
     } = options;
     const mode = useRef<MutationMode>(mutationMode);
     const paramsRef = useRef<Partial<UpdateParams<RecordType>>>(params);
     const snapshot = useRef<Snapshot>([]);
+    // Ref that stores the mutation with middlewares to avoid losing them if the calling component is unmounted
+    const mutateWithMiddlewares = useRef(dataProvider.update);
     // We need to store the call-time onError and onSettled in refs to be able to call them in the useMutation hook even
     // when the calling component is unmounted
     const callTimeOnError =
@@ -209,20 +211,9 @@ export const useUpdate = <RecordType extends RaRecord = any, ErrorType = Error>(
                     'useUpdate mutation requires a non-empty data object'
                 );
             }
-            if (mutateWithMiddlewares) {
-                return mutateWithMiddlewares(
-                    dataProvider.update.bind(dataProvider),
-                    callTimeResource,
-                    {
-                        id: callTimeId,
-                        data: callTimeData,
-                        previousData: callTimePreviousData,
-                        meta: callTimeMeta,
-                    }
-                ).then(({ data }) => data);
-            }
-            return dataProvider
-                .update<RecordType>(callTimeResource, {
+
+            return mutateWithMiddlewares
+                .current(callTimeResource, {
                     id: callTimeId,
                     data: callTimeData,
                     previousData: callTimePreviousData,
@@ -339,6 +330,15 @@ export const useUpdate = <RecordType extends RaRecord = any, ErrorType = Error>(
             ...otherCallTimeOptions
         } = callTimeOptions;
 
+        // Store the mutation with middlewares to avoid losing them if the calling component is unmounted
+        if (getMutateWithMiddlewares) {
+            mutateWithMiddlewares.current = getMutateWithMiddlewares(
+                dataProvider.update.bind(dataProvider)
+            );
+        } else {
+            mutateWithMiddlewares.current = dataProvider.update;
+        }
+
         // We need to keep the onSuccess callback here and not in the useMutation for undoable mutations
         hasCallTimeOnSuccess.current = !!onSuccess;
         // We need to store the onError and onSettled callbacks here to be able to call them in the useMutation hook
@@ -515,11 +515,12 @@ export type UseUpdateOptions<
 > & {
     mutationMode?: MutationMode;
     returnPromise?: boolean;
-    mutateWithMiddlewares?: <
+    getMutateWithMiddlewares?: <
         UpdateFunctionType extends
             DataProvider['update'] = DataProvider['update'],
     >(
-        mutate: UpdateFunctionType,
+        mutate: UpdateFunctionType
+    ) => (
         ...Params: Parameters<UpdateFunctionType>
     ) => ReturnType<UpdateFunctionType>;
 };
diff --git a/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx b/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx
index 382c3d4121b..74c43c2122f 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx
@@ -248,7 +248,7 @@ const WithMiddlewaresCore = () => {
         },
         {
             mutationMode: 'undoable',
-            mutateWithMiddlewares: async (mutate, resource, params) => {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
                     data: { title: `${params.data.title} from middleware` },
@@ -356,7 +356,7 @@ const WithMiddlewaresErrorCore = () => {
         },
         {
             mutationMode: 'undoable',
-            mutateWithMiddlewares: async (mutate, resource, params) => {
+            getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
                     data: { title: `${params.data.title} from middleware` },

From b506e6b89089d8b697e90f39904911fe45b8e560 Mon Sep 17 00:00:00 2001
From: Gildas Garcia <1122076+djhi@users.noreply.github.com>
Date: Mon, 27 May 2024 11:09:29 +0200
Subject: [PATCH 2/3] Add useRegisterMutationMiddleware tests

---
 .../useRegisterMutationMiddleware.spec.tsx    | 102 ++++++++++++++++++
 1 file changed, 102 insertions(+)
 create mode 100644 packages/ra-core/src/controller/saveContext/useRegisterMutationMiddleware.spec.tsx

diff --git a/packages/ra-core/src/controller/saveContext/useRegisterMutationMiddleware.spec.tsx b/packages/ra-core/src/controller/saveContext/useRegisterMutationMiddleware.spec.tsx
new file mode 100644
index 00000000000..6592a7ba0dc
--- /dev/null
+++ b/packages/ra-core/src/controller/saveContext/useRegisterMutationMiddleware.spec.tsx
@@ -0,0 +1,102 @@
+import React from 'react';
+import { fireEvent, render, screen, waitFor } from '@testing-library/react';
+import {
+    SaveContextProvider,
+    useMutationMiddlewares,
+    useRegisterMutationMiddleware,
+} from '../..';
+
+describe('useRegisterMutationMiddleware', () => {
+    it('should register and unregister middlewares correctly', async () => {
+        const middleware = jest.fn((next: () => void) => next());
+        const save = jest.fn();
+        const Middleware = () => {
+            useRegisterMutationMiddleware(middleware);
+            return null;
+        };
+        const Parent = () => {
+            const [mount, setMount] = React.useState(true);
+            const middlewaresFunctions = useMutationMiddlewares();
+
+            return (
+                <SaveContextProvider value={middlewaresFunctions}>
+                    {mount && <Middleware />}
+                    <button onClick={() => setMount(!mount)}>
+                        Toggle middleware
+                    </button>
+                    <button
+                        onClick={() => {
+                            const saveWithMiddlewares =
+                                middlewaresFunctions.getMutateWithMiddlewares(
+                                    save
+                                );
+                            saveWithMiddlewares();
+                        }}
+                    >
+                        Save
+                    </button>
+                </SaveContextProvider>
+            );
+        };
+
+        render(<Parent />);
+        fireEvent.click(screen.getByText('Save'));
+        await waitFor(() => {
+            expect(save).toHaveBeenCalledTimes(1);
+        });
+        expect(middleware).toHaveBeenCalledTimes(1);
+        save.mockClear();
+        middleware.mockClear();
+        fireEvent.click(screen.getByText('Toggle middleware'));
+        fireEvent.click(screen.getByText('Save'));
+        await waitFor(() => {
+            expect(save).toHaveBeenCalledTimes(1);
+        });
+        expect(middleware).not.toHaveBeenCalled();
+    });
+
+    it('should execute middlewares registered even if they have been unregistered as an optimistic side effect', async () => {
+        const middleware = jest.fn((next: () => void) => next());
+        const save = jest.fn();
+        const Middleware = () => {
+            useRegisterMutationMiddleware(middleware);
+            return <span>Middleware</span>;
+        };
+        const Parent = () => {
+            const [mount, setMount] = React.useState(true);
+            const middlewaresFunctions = useMutationMiddlewares();
+
+            return (
+                <SaveContextProvider value={middlewaresFunctions}>
+                    {mount && <Middleware />}
+                    <button
+                        onClick={() => {
+                            const saveWithMiddlewares =
+                                middlewaresFunctions.getMutateWithMiddlewares(
+                                    save
+                                );
+                            // Mimic optimistic side effect such as redirect which would unregister the middleware
+                            setMount(false);
+                            setTimeout(() => {
+                                saveWithMiddlewares();
+                            }, 250);
+                        }}
+                    >
+                        Save
+                    </button>
+                </SaveContextProvider>
+            );
+        };
+
+        render(<Parent />);
+        fireEvent.click(screen.getByText('Save'));
+        await waitFor(() => {
+            expect(screen.queryByText('Middleware')).toBeNull();
+            expect(save).not.toHaveBeenCalled();
+        });
+        await waitFor(() => {
+            expect(save).toHaveBeenCalledTimes(1);
+        });
+        expect(middleware).toHaveBeenCalledTimes(1);
+    });
+});

From 4e24ff43c60fb14d711fbf93a873a45a24038238 Mon Sep 17 00:00:00 2001
From: Gildas Garcia <1122076+djhi@users.noreply.github.com>
Date: Mon, 27 May 2024 11:34:46 +0200
Subject: [PATCH 3/3] Silence TS errors in stories

---
 packages/ra-core/src/dataProvider/useCreate.stories.tsx         | 2 ++
 .../ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx   | 2 ++
 .../ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx  | 2 ++
 .../ra-core/src/dataProvider/useUpdate.undoable.stories.tsx     | 2 ++
 4 files changed, 8 insertions(+)

diff --git a/packages/ra-core/src/dataProvider/useCreate.stories.tsx b/packages/ra-core/src/dataProvider/useCreate.stories.tsx
index 3d97dc26729..41b0de4fb32 100644
--- a/packages/ra-core/src/dataProvider/useCreate.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useCreate.stories.tsx
@@ -186,6 +186,7 @@ const WithMiddlewaresSuccessCore = () => {
             data: { title: 'Hello World' },
         },
         {
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
@@ -265,6 +266,7 @@ const WithMiddlewaresErrorCore = () => {
             data: { title: 'Hello World' },
         },
         {
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
diff --git a/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx b/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx
index fa92f99626c..d437c3ceef5 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useUpdate.optimistic.stories.tsx
@@ -190,6 +190,7 @@ const WithMiddlewaresSuccessCore = () => {
         },
         {
             mutationMode: 'optimistic',
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
@@ -270,6 +271,7 @@ const WithMiddlewaresErrorCore = () => {
         },
         {
             mutationMode: 'optimistic',
+            // @ts-ignore
             mutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
diff --git a/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx b/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx
index a86c5c4ce6c..50989984f89 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useUpdate.pessimistic.stories.tsx
@@ -188,6 +188,7 @@ const WithMiddlewaresSuccessCore = () => {
         },
         {
             mutationMode: 'pessimistic',
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
@@ -269,6 +270,7 @@ const WithMiddlewaresErrorCore = () => {
         },
         {
             mutationMode: 'pessimistic',
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
diff --git a/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx b/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx
index 74c43c2122f..c6e72089ad3 100644
--- a/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx
+++ b/packages/ra-core/src/dataProvider/useUpdate.undoable.stories.tsx
@@ -248,6 +248,7 @@ const WithMiddlewaresCore = () => {
         },
         {
             mutationMode: 'undoable',
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,
@@ -356,6 +357,7 @@ const WithMiddlewaresErrorCore = () => {
         },
         {
             mutationMode: 'undoable',
+            // @ts-ignore
             getMutateWithMiddlewares: mutate => async (resource, params) => {
                 return mutate(resource, {
                     ...params,