From a79bebd9fbdca8537f9dce2953e70e81514f461f Mon Sep 17 00:00:00 2001
From: Kasper Peulen <kasperpeulen@gmail.com>
Date: Thu, 2 May 2024 17:37:57 +0200
Subject: [PATCH] Wait for redirect to have happened

---
 .../nextjs/src/export-mocks/navigation/index.ts      |  4 ++--
 .../ServerActions.stories.tsx                        | 12 ++++++++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/code/frameworks/nextjs/src/export-mocks/navigation/index.ts b/code/frameworks/nextjs/src/export-mocks/navigation/index.ts
index 81a0ab6f20d2..f55ee86c36ca 100644
--- a/code/frameworks/nextjs/src/export-mocks/navigation/index.ts
+++ b/code/frameworks/nextjs/src/export-mocks/navigation/index.ts
@@ -59,13 +59,13 @@ export * from 'next/dist/client/components/navigation';
 
 // mock utilities/overrides (as of Next v14.2.0)
 export const redirect = fn(
-  (url: string, type: actual.RedirectType = actual.RedirectType.replace): never => {
+  (url: string, type: actual.RedirectType = actual.RedirectType.push): never => {
     throw getRedirectError(url, type, RedirectStatusCode.SeeOther);
   }
 ).mockName('next/navigation::redirect');
 
 export const permanentRedirect = fn(
-  (url: string, type: actual.RedirectType = actual.RedirectType.replace): never => {
+  (url: string, type: actual.RedirectType = actual.RedirectType.push): never => {
     throw getRedirectError(url, type, RedirectStatusCode.SeeOther);
   }
 ).mockName('next/navigation::permanentRedirect');
diff --git a/code/frameworks/nextjs/template/stories_nextjs-default-ts/ServerActions.stories.tsx b/code/frameworks/nextjs/template/stories_nextjs-default-ts/ServerActions.stories.tsx
index f330b2af3b9b..cb3f5bbdb2dd 100644
--- a/code/frameworks/nextjs/template/stories_nextjs-default-ts/ServerActions.stories.tsx
+++ b/code/frameworks/nextjs/template/stories_nextjs-default-ts/ServerActions.stories.tsx
@@ -1,9 +1,9 @@
 import React from 'react';
 import type { Meta, StoryObj } from '@storybook/react';
-import { expect, within, userEvent } from '@storybook/test';
+import { expect, within, userEvent, waitFor } from '@storybook/test';
 import { cookies } from '@storybook/nextjs/headers.mock';
 import { revalidatePath } from '@storybook/nextjs/cache.mock';
-import { redirect } from '@storybook/nextjs/navigation.mock';
+import { redirect, getRouter } from '@storybook/nextjs/navigation.mock';
 
 import { accessRoute, login, logout } from './server-actions';
 
@@ -63,6 +63,8 @@ export const ProtectedWhileLoggedOut: StoryObj<typeof Component> = {
 
     await expect(cookies().get).toHaveBeenCalledWith('user');
     await expect(redirect).toHaveBeenCalledWith('/');
+
+    await waitFor(() => expect(getRouter().push).toHaveBeenCalled());
   },
 };
 
@@ -77,6 +79,8 @@ export const ProtectedWhileLoggedIn: StoryObj<typeof Component> = {
     await expect(cookies().get).toHaveBeenLastCalledWith('user');
     await expect(revalidatePath).toHaveBeenLastCalledWith('/');
     await expect(redirect).toHaveBeenLastCalledWith('/protected');
+
+    await waitFor(() => expect(getRouter().push).toHaveBeenCalled());
   },
 };
 
@@ -91,6 +95,8 @@ export const Logout: StoryObj<typeof Component> = {
     await expect(cookies().delete).toHaveBeenCalled();
     await expect(revalidatePath).toHaveBeenCalledWith('/');
     await expect(redirect).toHaveBeenCalledWith('/');
+
+    await waitFor(() => expect(getRouter().push).toHaveBeenCalled());
   },
 };
 
@@ -102,5 +108,7 @@ export const Login: StoryObj<typeof Component> = {
     await expect(cookies().set).toHaveBeenCalledWith('user', 'storybookjs');
     await expect(revalidatePath).toHaveBeenCalledWith('/');
     await expect(redirect).toHaveBeenCalledWith('/');
+
+    await waitFor(() => expect(getRouter().push).toHaveBeenCalled());
   },
 };