diff --git a/app/note/edit/[id]/page.stories.tsx b/app/note/edit/[id]/page.stories.tsx
index 4c587c7..2a4b7be 100644
--- a/app/note/edit/[id]/page.stories.tsx
+++ b/app/note/edit/[id]/page.stories.tsx
@@ -1,4 +1,4 @@
-import { expect, fireEvent, userEvent, within } from '@storybook/test'
+import { expect, fireEvent, userEvent, waitFor, within } from '@storybook/test'
 import { Meta, StoryObj } from '@storybook/react'
 import { cookies } from '@storybook/nextjs/headers.mock'
 import Page from './page'
@@ -84,7 +84,7 @@ export const Save: Story = {
       await userEvent.click(
         await canvas.findByRole('menuitem', { name: /delete/i }),
       )
-      await expect(deleteNote).toHaveBeenCalledOnce()
+      await waitFor(() => expect(deleteNote).toHaveBeenCalledOnce())
       await expect(deleteNote).toHaveBeenCalledWith(2)
     })
   },
diff --git a/components/note-ui.stories.tsx b/components/note-ui.stories.tsx
index 444bec9..9c18b43 100644
--- a/components/note-ui.stories.tsx
+++ b/components/note-ui.stories.tsx
@@ -52,14 +52,14 @@ export const EditModeFlow: Story = {
       const saveButton = await canvas.findByRole('menuitem', { name: /done/i })
       await userEvent.click(saveButton)
       await expect(saveNote).toHaveBeenCalledOnce()
-      await expect(saveNote).toHaveBeenCalledWith('1', 'Edited Title', 'Edited Body')
+      await expect(saveNote).toHaveBeenCalledWith(1, 'Edited Title', 'Edited Body')
     })
 
     await step('Delete flow', async () => {
       const deleteButton = await canvas.findByRole('menuitem', { name: /delete/i })
       await userEvent.click(deleteButton)
       await expect(deleteNote).toHaveBeenCalledOnce()
-      await expect(deleteNote).toHaveBeenCalledWith('1')
+      await expect(deleteNote).toHaveBeenCalledWith(1)
     })
   },
 }