[notifications] Fix: #4959 transition on close

packages/toolpad-core/src/useNotifications/NotificationsProvider.tsx

@@ -55,8 +55,7 @@ interface NotificationProps {
 function Notification({ notificationKey, open, message, options, badge }: NotificationProps) {
   const globalLocaleText = useLocaleText();
   const localeText = { ...defaultLocaleText, ...globalLocaleText };
-  const { close } = useNonNullableContext(NotificationsContext);
-
+  const { close, remove } = useNonNullableContext(NotificationsContext);
   const { severity, actionText, onAction, autoHideDuration } = options;
 
   const handleClose = React.useCallback(
@@ -69,6 +68,10 @@ function Notification({ notificationKey, open, message, options, badge }: Notifi
     [notificationKey, close],
   );
 
+  const handleExited = React.useCallback(() => {
+    remove(notificationKey);
+  }, [notificationKey, remove]);
+
   const action = (
     <React.Fragment>
       {onAction ? (
@@ -90,10 +93,25 @@ function Notification({ notificationKey, open, message, options, badge }: Notifi
 
   const props = React.useContext(RootPropsContext);
   const SnackbarComponent = props?.slots?.snackbar ?? Snackbar;
+
+  // Passing `onExited` through `externalSlotProps` here.
+  // Passing it through `additionalProps` causes it to be overwritten when
+  // transition slotProps are specified in RootPropsContext.
+  const externalSnackbarSlotProps = props?.slotProps?.snackbar?.slotProps;
+  const externalTransitionProps = externalSnackbarSlotProps?.transition;
   const snackbarSlotProps = useSlotProps({
     elementType: SnackbarComponent,
     ownerState: props,
-    externalSlotProps: props?.slotProps?.snackbar,
+    externalSlotProps: {
+      ...props?.slotProps?.snackbar,
+      slotProps: {
+        ...externalSnackbarSlotProps,
+        transition: {
+          ...externalTransitionProps,
+          onExited: handleExited,
+        },
+      },
+    },
     additionalProps: {
       open,
       autoHideDuration,
@@ -158,6 +176,15 @@ const generateId = () => {
   return id;
 };
 
+export interface RemoveNotification {
+  /**
+   * Remove a snackbar from the application state (after it has been closed).
+   *
+   * @param key The key of the notification to remove.
+   */
+  (key: string): void;
+}
+
 /**
  * Provider for Notifications. The subtree of this component can use the `useNotifications` hook to
  * access the notifications API. The notifications are shown in the same order they are requested.
@@ -193,11 +220,18 @@ function NotificationsProvider(props: NotificationsProviderProps) {
   const close = React.useCallback<CloseNotification>((key) => {
     setState((prev) => ({
       ...prev,
-      queue: prev.queue.filter((n) => n.notificationKey !== key),
+      queue: prev.queue.map((n) => (n.notificationKey === key ? { ...n, open: false } : n)),
+    }));
+  }, []);
+
+  const remove = React.useCallback<RemoveNotification>((key) => {
+    setState((prev) => ({
+      ...prev,
+      queue: prev.queue.filter((n) => key !== n.notificationKey),
     }));
   }, []);
 
-  const contextValue = React.useMemo(() => ({ show, close }), [show, close]);
+  const contextValue = React.useMemo(() => ({ show, close, remove }), [show, close, remove]);
 
   return (
     <RootPropsContext.Provider value={props}>

packages/toolpad-core/src/useNotifications/useNotifications.test.tsx

@@ -4,7 +4,7 @@
 
 import * as React from 'react';
 import { describe, test, expect } from 'vitest';
-import { renderHook, within, screen } from '@testing-library/react';
+import { renderHook, within, screen, waitFor } from '@testing-library/react';
 import { userEvent } from '@testing-library/user-event';
 import { useNotifications } from './useNotifications';
 import { NotificationsProvider } from './NotificationsProvider';
@@ -35,6 +35,8 @@ describe('useNotifications', () => {
 
     rerender();
 
-    expect(screen.queryByRole('alert')).toBeNull();
+    await waitFor(() => {
+      expect(screen.queryByRole('alert')).toBeNull();
+    });
   });
 });