diff --git a/src/components/impersonation.spec.tsx b/src/components/impersonation.spec.tsx index 955ed5d..307f36b 100644 --- a/src/components/impersonation.spec.tsx +++ b/src/components/impersonation.spec.tsx @@ -129,6 +129,21 @@ describe('Impersonation', () => { render(); const stopButton = await screen.findByText('Stop'); stopButton.click(); - expect(handleSignOutAction).toHaveBeenCalled(); + expect(handleSignOutAction).toHaveBeenCalledWith({}); + }); + + it('should pass returnTo prop to handleSignOutAction when provided', async () => { + (useAuth as jest.Mock).mockReturnValue({ + impersonator: { email: 'admin@example.com' }, + user: { id: '123', email: 'user@example.com' }, + organizationId: null, + loading: false, + }); + + const returnTo = '/dashboard'; + render(); + const stopButton = await screen.findByText('Stop'); + stopButton.click(); + expect(handleSignOutAction).toHaveBeenCalledWith({ returnTo }); }); }); diff --git a/src/components/impersonation.tsx b/src/components/impersonation.tsx index a9c773e..306cae9 100644 --- a/src/components/impersonation.tsx +++ b/src/components/impersonation.tsx @@ -9,9 +9,10 @@ import { useAuth } from './authkit-provider.js'; interface ImpersonationProps extends React.ComponentPropsWithoutRef<'div'> { side?: 'top' | 'bottom'; + returnTo?: string; } -export function Impersonation({ side = 'bottom', ...props }: ImpersonationProps) { +export function Impersonation({ side = 'bottom', returnTo, ...props }: ImpersonationProps) { const { user, impersonator, organizationId, loading } = useAuth(); const [organization, setOrganization] = React.useState(null); @@ -78,7 +79,7 @@ export function Impersonation({ side = 'bottom', ...props }: ImpersonationProps)
{ event.preventDefault(); - await handleSignOutAction(); + await handleSignOutAction({ returnTo }); }} style={{ display: 'flex',