From 8e561dced401b247adbf8c609149a8f52e938741 Mon Sep 17 00:00:00 2001
From: Jeremy <git@jrmyphlmn.com>
Date: Tue, 11 Mar 2025 14:39:36 -0700
Subject: [PATCH] feat: add artifact tests (#859)

---
 components/artifact-close-button.tsx |   1 +
 components/artifact.tsx              |   1 +
 components/message-actions.tsx       |   4 +-
 components/multimodal-input.tsx      |   2 +-
 lib/ai/models.test.ts                | 218 +-----------------------
 playwright.config.ts                 |   9 +
 tests/artifacts.test.ts              |  74 ++++++++
 tests/chat.test.ts                   |  69 ++++++--
 tests/pages/artifact.ts              | 108 ++++++++++++
 tests/pages/chat.ts                  |  37 ++--
 tests/prompts/basic.ts               | 141 ++++++++++++++++
 tests/prompts/utils.ts               | 242 +++++++++++++++++++++++++++
 tests/reasoning.test.ts              |  18 +-
 13 files changed, 669 insertions(+), 255 deletions(-)
 create mode 100644 tests/artifacts.test.ts
 create mode 100644 tests/pages/artifact.ts
 create mode 100644 tests/prompts/basic.ts
 create mode 100644 tests/prompts/utils.ts

diff --git a/components/artifact-close-button.tsx b/components/artifact-close-button.tsx
index 89b899e..27fa9cd 100644
--- a/components/artifact-close-button.tsx
+++ b/components/artifact-close-button.tsx
@@ -8,6 +8,7 @@ function PureArtifactCloseButton() {
 
   return (
     <Button
+      data-testid="artifact-close-button"
       variant="outline"
       className="h-fit p-2 dark:hover:bg-zinc-700"
       onClick={() => {
diff --git a/components/artifact.tsx b/components/artifact.tsx
index f031c03..5be9f1a 100644
--- a/components/artifact.tsx
+++ b/components/artifact.tsx
@@ -269,6 +269,7 @@ function PureArtifact({
     <AnimatePresence>
       {artifact.isVisible && (
         <motion.div
+          data-testid="artifact"
           className="flex flex-row h-dvh w-dvw fixed top-0 left-0 z-50 bg-transparent"
           initial={{ opacity: 1 }}
           animate={{ opacity: 1 }}
diff --git a/components/message-actions.tsx b/components/message-actions.tsx
index dc67dc5..7ec9536 100644
--- a/components/message-actions.tsx
+++ b/components/message-actions.tsx
@@ -1,5 +1,4 @@
 import type { Message } from 'ai';
-import { toast } from 'sonner';
 import { useSWRConfig } from 'swr';
 import { useCopyToClipboard } from 'usehooks-ts';
 
@@ -15,6 +14,7 @@ import {
 } from './ui/tooltip';
 import { memo } from 'react';
 import equal from 'fast-deep-equal';
+import { toast } from 'sonner';
 
 export function PureMessageActions({
   chatId,
@@ -57,6 +57,7 @@ export function PureMessageActions({
         <Tooltip>
           <TooltipTrigger asChild>
             <Button
+              data-testid="message-upvote"
               className="py-1 px-2 h-fit text-muted-foreground !pointer-events-auto"
               disabled={vote?.isUpvoted}
               variant="outline"
@@ -109,6 +110,7 @@ export function PureMessageActions({
         <Tooltip>
           <TooltipTrigger asChild>
             <Button
+              data-testid="message-downvote"
               className="py-1 px-2 h-fit text-muted-foreground !pointer-events-auto"
               variant="outline"
               disabled={vote && !vote.isUpvoted}
diff --git a/components/multimodal-input.tsx b/components/multimodal-input.tsx
index 2e5475b..2eeb93c 100644
--- a/components/multimodal-input.tsx
+++ b/components/multimodal-input.tsx
@@ -247,7 +247,7 @@ function PureMultimodalInput({
         autoFocus
         onKeyDown={(event) => {
           if (
-            event.key === "Enter" &&
+            event.key === 'Enter' &&
             !event.shiftKey &&
             !event.nativeEvent.isComposing
           ) {
diff --git a/lib/ai/models.test.ts b/lib/ai/models.test.ts
index f39fcd9..a487134 100644
--- a/lib/ai/models.test.ts
+++ b/lib/ai/models.test.ts
@@ -1,180 +1,6 @@
-import { CoreMessage, FinishReason, simulateReadableStream } from 'ai';
+import { simulateReadableStream } from 'ai';
 import { MockLanguageModelV1 } from 'ai/test';
-
-interface ReasoningChunk {
-  type: 'reasoning';
-  textDelta: string;
-}
-
-interface TextDeltaChunk {
-  type: 'text-delta';
-  textDelta: string;
-}
-
-interface FinishChunk {
-  type: 'finish';
-  finishReason: FinishReason;
-  logprobs: undefined;
-  usage: { completionTokens: number; promptTokens: number };
-}
-
-type Chunk = TextDeltaChunk | ReasoningChunk | FinishChunk;
-
-const getResponseChunksByPrompt = (
-  prompt: CoreMessage[],
-  isReasoningEnabled: boolean = false,
-): Array<Chunk> => {
-  const userMessage = prompt.at(-1);
-
-  if (!userMessage) {
-    throw new Error('No user message found');
-  }
-
-  if (isReasoningEnabled) {
-    if (
-      compareMessages(userMessage, {
-        role: 'user',
-        content: [{ type: 'text', text: 'why is the sky blue?' }],
-      })
-    ) {
-      return [
-        { type: 'reasoning', textDelta: 'the ' },
-        { type: 'reasoning', textDelta: 'sky ' },
-        { type: 'reasoning', textDelta: 'is ' },
-        { type: 'reasoning', textDelta: 'blue ' },
-        { type: 'reasoning', textDelta: 'because ' },
-        { type: 'reasoning', textDelta: 'of ' },
-        { type: 'reasoning', textDelta: 'rayleigh ' },
-        { type: 'reasoning', textDelta: 'scattering! ' },
-        { type: 'text-delta', textDelta: "it's " },
-        { type: 'text-delta', textDelta: 'just ' },
-        { type: 'text-delta', textDelta: 'blue ' },
-        { type: 'text-delta', textDelta: 'duh!' },
-        {
-          type: 'finish',
-          finishReason: 'stop',
-          logprobs: undefined,
-          usage: { completionTokens: 10, promptTokens: 3 },
-        },
-      ];
-    } else if (
-      compareMessages(userMessage, {
-        role: 'user',
-        content: [{ type: 'text', text: 'why is grass green?' }],
-      })
-    ) {
-      return [
-        { type: 'reasoning', textDelta: 'grass ' },
-        { type: 'reasoning', textDelta: 'is ' },
-        { type: 'reasoning', textDelta: 'green ' },
-        { type: 'reasoning', textDelta: 'because ' },
-        { type: 'reasoning', textDelta: 'of ' },
-        { type: 'reasoning', textDelta: 'chlorophyll ' },
-        { type: 'reasoning', textDelta: 'absorption! ' },
-        { type: 'text-delta', textDelta: "it's " },
-        { type: 'text-delta', textDelta: 'just ' },
-        { type: 'text-delta', textDelta: 'green ' },
-        { type: 'text-delta', textDelta: 'duh!' },
-        {
-          type: 'finish',
-          finishReason: 'stop',
-          logprobs: undefined,
-          usage: { completionTokens: 10, promptTokens: 3 },
-        },
-      ];
-    }
-  }
-
-  if (
-    compareMessages(userMessage, {
-      role: 'user',
-      content: [{ type: 'text', text: 'why is grass green?' }],
-    })
-  ) {
-    return [
-      { type: 'text-delta', textDelta: "it's " },
-      { type: 'text-delta', textDelta: 'just ' },
-      { type: 'text-delta', textDelta: 'green ' },
-      { type: 'text-delta', textDelta: 'duh!' },
-      {
-        type: 'finish',
-        finishReason: 'stop',
-        logprobs: undefined,
-        usage: { completionTokens: 10, promptTokens: 3 },
-      },
-    ];
-  } else if (
-    compareMessages(userMessage, {
-      role: 'user',
-      content: [{ type: 'text', text: 'why is the sky blue?' }],
-    })
-  ) {
-    return [
-      { type: 'text-delta', textDelta: "it's " },
-      { type: 'text-delta', textDelta: 'just ' },
-      { type: 'text-delta', textDelta: 'blue ' },
-      { type: 'text-delta', textDelta: 'duh!' },
-      {
-        type: 'finish',
-        finishReason: 'stop',
-        logprobs: undefined,
-        usage: { completionTokens: 10, promptTokens: 3 },
-      },
-    ];
-  } else if (
-    compareMessages(userMessage, {
-      role: 'user',
-      content: [
-        { type: 'text', text: 'What are the advantages of using Next.js?' },
-      ],
-    })
-  ) {
-    return [
-      { type: 'text-delta', textDelta: 'with ' },
-      { type: 'text-delta', textDelta: 'next.js ' },
-      { type: 'text-delta', textDelta: 'you ' },
-      { type: 'text-delta', textDelta: 'can ' },
-      { type: 'text-delta', textDelta: 'ship ' },
-      { type: 'text-delta', textDelta: 'fast! ' },
-      {
-        type: 'finish',
-        finishReason: 'stop',
-        logprobs: undefined,
-        usage: { completionTokens: 10, promptTokens: 3 },
-      },
-    ];
-  } else if (
-    compareMessages(userMessage, {
-      role: 'user',
-      content: [
-        {
-          type: 'text',
-          text: 'who painted this?',
-        },
-        {
-          type: 'image',
-          image: '...',
-        },
-      ],
-    })
-  ) {
-    return [
-      { type: 'text-delta', textDelta: 'this ' },
-      { type: 'text-delta', textDelta: 'painting ' },
-      { type: 'text-delta', textDelta: 'is ' },
-      { type: 'text-delta', textDelta: 'by ' },
-      { type: 'text-delta', textDelta: 'monet!' },
-      {
-        type: 'finish',
-        finishReason: 'stop',
-        logprobs: undefined,
-        usage: { completionTokens: 10, promptTokens: 3 },
-      },
-    ];
-  }
-
-  return [];
-};
+import { getResponseChunksByPrompt } from '@/tests/prompts/utils';
 
 export const chatModel = new MockLanguageModelV1({
   doGenerate: async () => ({
@@ -236,46 +62,10 @@ export const artifactModel = new MockLanguageModelV1({
     usage: { promptTokens: 10, completionTokens: 20 },
     text: `Hello, world!`,
   }),
-  doStream: async () => ({
+  doStream: async ({ prompt }) => ({
     stream: simulateReadableStream({
-      chunks: [
-        { type: 'text-delta', textDelta: 'test' },
-        {
-          type: 'finish',
-          finishReason: 'stop',
-          logprobs: undefined,
-          usage: { completionTokens: 10, promptTokens: 3 },
-        },
-      ],
+      chunks: getResponseChunksByPrompt(prompt),
     }),
     rawCall: { rawPrompt: null, rawSettings: {} },
   }),
 });
-
-function compareMessages(msg1: CoreMessage, msg2: CoreMessage): boolean {
-  if (msg1.role !== msg2.role) return false;
-
-  if (!Array.isArray(msg1.content) || !Array.isArray(msg2.content)) {
-    return false;
-  }
-
-  if (msg1.content.length !== msg2.content.length) return false;
-
-  for (let i = 0; i < msg1.content.length; i++) {
-    const item1 = msg1.content[i];
-    const item2 = msg2.content[i];
-
-    if (item1.type !== item2.type) return false;
-
-    if (item1.type === 'image' && item2.type === 'image') {
-      // if (item1.image.toString() !== item2.image.toString()) return false;
-      // if (item1.mimeType !== item2.mimeType) return false;
-    } else if (item1.type === 'text' && item2.type === 'text') {
-      if (item1.text !== item2.text) return false;
-    } else {
-      return false;
-    }
-  }
-
-  return true;
-}
diff --git a/playwright.config.ts b/playwright.config.ts
index 89cbb5a..4b6499d 100644
--- a/playwright.config.ts
+++ b/playwright.config.ts
@@ -82,6 +82,15 @@ export default defineConfig({
         storageState: 'playwright/.reasoning/session.json',
       },
     },
+    {
+      name: 'artifacts',
+      testMatch: /artifacts.test.ts/,
+      dependencies: ['setup:auth'],
+      use: {
+        ...devices['Desktop Chrome'],
+        storageState: 'playwright/.auth/session.json',
+      },
+    },
 
     // {
     //   name: 'firefox',
diff --git a/tests/artifacts.test.ts b/tests/artifacts.test.ts
new file mode 100644
index 0000000..eb2fee6
--- /dev/null
+++ b/tests/artifacts.test.ts
@@ -0,0 +1,74 @@
+import { expect, test } from '@playwright/test';
+import { ChatPage } from './pages/chat';
+import { ArtifactPage } from './pages/artifact';
+
+test.describe('artifacts activity', () => {
+  let chatPage: ChatPage;
+  let artifactPage: ArtifactPage;
+
+  test.beforeEach(async ({ page }) => {
+    chatPage = new ChatPage(page);
+    artifactPage = new ArtifactPage(page);
+
+    await chatPage.createNewChat();
+  });
+
+  test('create a text artifact', async () => {
+    await chatPage.createNewChat();
+
+    await chatPage.sendUserMessage(
+      'Help me write an essay about Silicon Valley',
+    );
+    await artifactPage.isGenerationComplete();
+
+    expect(artifactPage.artifact).toBeVisible();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toBe(
+      'A document was created and is now visible to the user.',
+    );
+
+    await chatPage.hasChatIdInUrl();
+  });
+
+  test('toggle artifact visibility', async () => {
+    await chatPage.createNewChat();
+
+    await chatPage.sendUserMessage(
+      'Help me write an essay about Silicon Valley',
+    );
+    await artifactPage.isGenerationComplete();
+
+    expect(artifactPage.artifact).toBeVisible();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toBe(
+      'A document was created and is now visible to the user.',
+    );
+
+    await artifactPage.closeArtifact();
+    await chatPage.isElementNotVisible('artifact');
+  });
+
+  test('send follow up message after generation', async () => {
+    await chatPage.createNewChat();
+
+    await chatPage.sendUserMessage(
+      'Help me write an essay about Silicon Valley',
+    );
+    await artifactPage.isGenerationComplete();
+
+    expect(artifactPage.artifact).toBeVisible();
+
+    const assistantMessage = await artifactPage.getRecentAssistantMessage();
+    expect(assistantMessage.content).toBe(
+      'A document was created and is now visible to the user.',
+    );
+
+    await artifactPage.sendUserMessage('Thanks!');
+    await artifactPage.isGenerationComplete();
+
+    const secondAssistantMessage = await chatPage.getRecentAssistantMessage();
+    expect(secondAssistantMessage.content).toBe("You're welcome!");
+  });
+});
diff --git a/tests/chat.test.ts b/tests/chat.test.ts
index 7842a6f..869fa57 100644
--- a/tests/chat.test.ts
+++ b/tests/chat.test.ts
@@ -1,5 +1,5 @@
-import { test, expect } from '@playwright/test';
 import { ChatPage } from './pages/chat';
+import { test, expect } from '@playwright/test';
 
 test.describe('chat activity', () => {
   let chatPage: ChatPage;
@@ -10,19 +10,19 @@ test.describe('chat activity', () => {
   });
 
   test('send a user message and receive response', async () => {
-    await chatPage.sendUserMessage('why is grass green?');
+    await chatPage.sendUserMessage('Why is grass green?');
     await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
-    expect(assistantMessage.content).toContain("it's just green duh!");
+    expect(assistantMessage.content).toContain("It's just green duh!");
   });
 
   test('redirect to /chat/:id after submitting message', async () => {
-    await chatPage.sendUserMessage('why is grass green?');
+    await chatPage.sendUserMessage('Why is grass green?');
     await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
-    expect(assistantMessage.content).toContain("it's just green duh!");
+    expect(assistantMessage.content).toContain("It's just green duh!");
     await chatPage.hasChatIdInUrl();
   });
 
@@ -32,7 +32,7 @@ test.describe('chat activity', () => {
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
     expect(assistantMessage.content).toContain(
-      'with next.js you can ship fast!',
+      'With Next.js, you can ship fast!',
     );
   });
 
@@ -40,7 +40,7 @@ test.describe('chat activity', () => {
     await expect(chatPage.sendButton).toBeVisible();
     await expect(chatPage.sendButton).toBeDisabled();
 
-    await chatPage.sendUserMessage('why is grass green?');
+    await chatPage.sendUserMessage('Why is grass green?');
 
     await expect(chatPage.sendButton).not.toBeVisible();
     await expect(chatPage.stopButton).toBeVisible();
@@ -52,26 +52,26 @@ test.describe('chat activity', () => {
   });
 
   test('stop generation during submission', async () => {
-    await chatPage.sendUserMessage('why is grass green?');
+    await chatPage.sendUserMessage('Why is grass green?');
     await expect(chatPage.stopButton).toBeVisible();
     await chatPage.stopButton.click();
     await expect(chatPage.sendButton).toBeVisible();
   });
 
   test('edit user message and resubmit', async () => {
-    await chatPage.sendUserMessage('why is grass green?');
+    await chatPage.sendUserMessage('Why is grass green?');
     await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
-    expect(assistantMessage.content).toContain("it's just green duh!");
+    expect(assistantMessage.content).toContain("It's just green duh!");
 
     const userMessage = await chatPage.getRecentUserMessage();
-    await userMessage.edit('why is the sky blue?');
+    await userMessage.edit('Why is the sky blue?');
 
     await chatPage.isGenerationComplete();
 
     const updatedAssistantMessage = await chatPage.getRecentAssistantMessage();
-    expect(updatedAssistantMessage.content).toContain("it's just blue duh!");
+    expect(updatedAssistantMessage.content).toContain("It's just blue duh!");
   });
 
   test('hide suggested actions after sending message', async () => {
@@ -87,7 +87,7 @@ test.describe('chat activity', () => {
     await chatPage.isElementVisible('input-attachment-loader');
     await chatPage.isElementNotVisible('input-attachment-loader');
 
-    await chatPage.sendUserMessage('who painted this?');
+    await chatPage.sendUserMessage('Who painted this?');
 
     const userMessage = await chatPage.getRecentUserMessage();
     expect(userMessage.attachments).toHaveLength(1);
@@ -95,6 +95,47 @@ test.describe('chat activity', () => {
     await chatPage.isGenerationComplete();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
-    expect(assistantMessage.content).toBe('this painting is by monet!');
+    expect(assistantMessage.content).toBe('This painting is by Monet!');
+  });
+
+  test('call weather tool', async () => {
+    await chatPage.sendUserMessage("What's the weather in sf?");
+    await chatPage.isGenerationComplete();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+
+    expect(assistantMessage.content).toBe(
+      'The current temperature in San Francisco is 17°C.',
+    );
+  });
+
+  test('upvote message', async () => {
+    await chatPage.sendUserMessage('Why is the sky blue?');
+    await chatPage.isGenerationComplete();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    await assistantMessage.upvote();
+    await chatPage.isVoteComplete();
+  });
+
+  test('downvote message', async () => {
+    await chatPage.sendUserMessage('Why is the sky blue?');
+    await chatPage.isGenerationComplete();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    await assistantMessage.downvote();
+    await chatPage.isVoteComplete();
+  });
+
+  test('update vote', async () => {
+    await chatPage.sendUserMessage('Why is the sky blue?');
+    await chatPage.isGenerationComplete();
+
+    const assistantMessage = await chatPage.getRecentAssistantMessage();
+    await assistantMessage.upvote();
+    await chatPage.isVoteComplete();
+
+    await assistantMessage.downvote();
+    await chatPage.isVoteComplete();
   });
 });
diff --git a/tests/pages/artifact.ts b/tests/pages/artifact.ts
new file mode 100644
index 0000000..23cc263
--- /dev/null
+++ b/tests/pages/artifact.ts
@@ -0,0 +1,108 @@
+import { expect, Page } from '@playwright/test';
+
+export class ArtifactPage {
+  constructor(private page: Page) {}
+
+  public get artifact() {
+    return this.page.getByTestId('artifact');
+  }
+
+  public get sendButton() {
+    return this.artifact.getByTestId('send-button');
+  }
+
+  public get stopButton() {
+    return this.page.getByTestId('stop-button');
+  }
+
+  public get multimodalInput() {
+    return this.page.getByTestId('multimodal-input');
+  }
+
+  async isGenerationComplete() {
+    const response = await this.page.waitForResponse((response) =>
+      response.url().includes('/api/chat'),
+    );
+
+    await response.finished();
+  }
+
+  async sendUserMessage(message: string) {
+    await this.artifact.getByTestId('multimodal-input').click();
+    await this.artifact.getByTestId('multimodal-input').fill(message);
+    await this.artifact.getByTestId('send-button').click();
+  }
+
+  async getRecentAssistantMessage() {
+    const messageElements = await this.artifact
+      .getByTestId('message-assistant')
+      .all();
+    const lastMessageElement = messageElements[messageElements.length - 1];
+
+    const content = await lastMessageElement
+      .getByTestId('message-content')
+      .innerText()
+      .catch(() => null);
+
+    const reasoningElement = await lastMessageElement
+      .getByTestId('message-reasoning')
+      .isVisible()
+      .then(async (visible) =>
+        visible
+          ? await lastMessageElement
+              .getByTestId('message-reasoning')
+              .innerText()
+          : null,
+      )
+      .catch(() => null);
+
+    return {
+      element: lastMessageElement,
+      content,
+      reasoning: reasoningElement,
+      async toggleReasoningVisibility() {
+        await lastMessageElement
+          .getByTestId('message-reasoning-toggle')
+          .click();
+      },
+    };
+  }
+
+  async getRecentUserMessage() {
+    const messageElements = await this.artifact
+      .getByTestId('message-user')
+      .all();
+    const lastMessageElement = messageElements[messageElements.length - 1];
+
+    const content = await lastMessageElement.innerText();
+
+    const hasAttachments = await lastMessageElement
+      .getByTestId('message-attachments')
+      .isVisible()
+      .catch(() => false);
+
+    const attachments = hasAttachments
+      ? await lastMessageElement.getByTestId('message-attachments').all()
+      : [];
+
+    const page = this.artifact;
+
+    return {
+      element: lastMessageElement,
+      content,
+      attachments,
+      async edit(newMessage: string) {
+        await page.getByTestId('message-edit').click();
+        await page.getByTestId('message-editor').fill(newMessage);
+        await page.getByTestId('message-editor-send-button').click();
+        await expect(
+          page.getByTestId('message-editor-send-button'),
+        ).not.toBeVisible();
+      },
+    };
+  }
+
+  async closeArtifact() {
+    return this.page.getByTestId('artifact-close-button').click();
+  }
+}
diff --git a/tests/pages/chat.ts b/tests/pages/chat.ts
index 0d4cdf8..7a99875 100644
--- a/tests/pages/chat.ts
+++ b/tests/pages/chat.ts
@@ -27,18 +27,25 @@ export class ChatPage {
   }
 
   async sendUserMessage(message: string) {
-    await this.page.getByTestId('multimodal-input').click();
-    await this.page.getByTestId('multimodal-input').fill(message);
-    await this.page.getByTestId('send-button').click();
-    await this.page.getByTestId('message-user').isVisible();
-    expect(await this.page.getByTestId('message-user').innerText()).toContain(
-      message,
-    );
+    await this.multimodalInput.click();
+    await this.multimodalInput.fill(message);
+    await this.sendButton.click();
   }
 
   async isGenerationComplete() {
-    await expect(this.page.getByTestId('send-button')).toBeVisible();
-    await this.page.waitForTimeout(2000);
+    const response = await this.page.waitForResponse((response) =>
+      response.url().includes('/api/chat'),
+    );
+
+    await response.finished();
+  }
+
+  async isVoteComplete() {
+    const response = await this.page.waitForResponse((response) =>
+      response.url().includes('/api/vote'),
+    );
+
+    await response.finished();
   }
 
   async hasChatIdInUrl() {
@@ -123,21 +130,21 @@ export class ChatPage {
       )
       .catch(() => null);
 
-    const page = this.page;
-
     return {
       element: lastMessageElement,
       content,
       reasoning: reasoningElement,
-      async waitForGenerationComplete() {
-        await expect(page.getByTestId('send-button')).toBeVisible();
-        await page.waitForTimeout(2000);
-      },
       async toggleReasoningVisibility() {
         await lastMessageElement
           .getByTestId('message-reasoning-toggle')
           .click();
       },
+      async upvote() {
+        await lastMessageElement.getByTestId('message-upvote').click();
+      },
+      async downvote() {
+        await lastMessageElement.getByTestId('message-downvote').click();
+      },
     };
   }
 
diff --git a/tests/prompts/basic.ts b/tests/prompts/basic.ts
new file mode 100644
index 0000000..4835e33
--- /dev/null
+++ b/tests/prompts/basic.ts
@@ -0,0 +1,141 @@
+import { CoreMessage } from 'ai';
+
+export const TEST_PROMPTS: Record<string, CoreMessage> = {
+  USER_SKY: {
+    role: 'user',
+    content: [{ type: 'text', text: 'Why is the sky blue?' }],
+  },
+  USER_GRASS: {
+    role: 'user',
+    content: [{ type: 'text', text: 'Why is grass green?' }],
+  },
+  USER_THANKS: {
+    role: 'user',
+    content: [{ type: 'text', text: 'Thanks!' }],
+  },
+  USER_NEXTJS: {
+    role: 'user',
+    content: [
+      { type: 'text', text: 'What are the advantages of using Next.js?' },
+    ],
+  },
+  USER_IMAGE_ATTACHMENT: {
+    role: 'user',
+    content: [
+      {
+        type: 'text',
+        text: 'Who painted this?',
+      },
+      {
+        type: 'image',
+        image: '...',
+      },
+    ],
+  },
+  USER_TEXT_ARTIFACT: {
+    role: 'user',
+    content: [
+      {
+        type: 'text',
+        text: 'Help me write an essay about Silicon Valley',
+      },
+    ],
+  },
+  CREATE_DOCUMENT_TEXT_CALL: {
+    role: 'user',
+    content: [
+      {
+        type: 'text',
+        text: 'Essay about Silicon Valley',
+      },
+    ],
+  },
+  CREATE_DOCUMENT_TEXT_RESULT: {
+    role: 'tool',
+    content: [
+      {
+        type: 'tool-result',
+        toolCallId: 'call_123',
+        toolName: 'createDocument',
+        result: {
+          id: '3ca386a4-40c6-4630-8ed1-84cbd46cc7eb',
+          title: 'Essay about Silicon Valley',
+          kind: 'text',
+          content: 'A document was created and is now visible to the user.',
+        },
+      },
+    ],
+  },
+  GET_WEATHER_CALL: {
+    role: 'user',
+    content: [
+      {
+        type: 'text',
+        text: "What's the weather in sf?",
+      },
+    ],
+  },
+  GET_WEATHER_RESULT: {
+    role: 'tool',
+    content: [
+      {
+        type: 'tool-result',
+        toolCallId: 'call_456',
+        toolName: 'getWeather',
+        result: {
+          latitude: 37.763283,
+          longitude: -122.41286,
+          generationtime_ms: 0.06449222564697266,
+          utc_offset_seconds: -25200,
+          timezone: 'America/Los_Angeles',
+          timezone_abbreviation: 'GMT-7',
+          elevation: 18,
+          current_units: {
+            time: 'iso8601',
+            interval: 'seconds',
+            temperature_2m: '°C',
+          },
+          current: {
+            time: '2025-03-10T14:00',
+            interval: 900,
+            temperature_2m: 17,
+          },
+          daily_units: {
+            time: 'iso8601',
+            sunrise: 'iso8601',
+            sunset: 'iso8601',
+          },
+          daily: {
+            time: [
+              '2025-03-10',
+              '2025-03-11',
+              '2025-03-12',
+              '2025-03-13',
+              '2025-03-14',
+              '2025-03-15',
+              '2025-03-16',
+            ],
+            sunrise: [
+              '2025-03-10T07:27',
+              '2025-03-11T07:25',
+              '2025-03-12T07:24',
+              '2025-03-13T07:22',
+              '2025-03-14T07:21',
+              '2025-03-15T07:19',
+              '2025-03-16T07:18',
+            ],
+            sunset: [
+              '2025-03-10T19:12',
+              '2025-03-11T19:13',
+              '2025-03-12T19:14',
+              '2025-03-13T19:15',
+              '2025-03-14T19:16',
+              '2025-03-15T19:17',
+              '2025-03-16T19:17',
+            ],
+          },
+        },
+      },
+    ],
+  },
+};
diff --git a/tests/prompts/utils.ts b/tests/prompts/utils.ts
new file mode 100644
index 0000000..a67618b
--- /dev/null
+++ b/tests/prompts/utils.ts
@@ -0,0 +1,242 @@
+import { CoreMessage, LanguageModelV1StreamPart } from 'ai';
+import { TEST_PROMPTS } from './basic';
+
+export function compareMessages(
+  firstMessage: CoreMessage,
+  secondMessage: CoreMessage,
+): boolean {
+  if (firstMessage.role !== secondMessage.role) return false;
+
+  if (
+    !Array.isArray(firstMessage.content) ||
+    !Array.isArray(secondMessage.content)
+  ) {
+    return false;
+  }
+
+  if (firstMessage.content.length !== secondMessage.content.length) {
+    return false;
+  }
+
+  for (let i = 0; i < firstMessage.content.length; i++) {
+    const item1 = firstMessage.content[i];
+    const item2 = secondMessage.content[i];
+
+    if (item1.type !== item2.type) return false;
+
+    if (item1.type === 'image' && item2.type === 'image') {
+      // if (item1.image.toString() !== item2.image.toString()) return false;
+      // if (item1.mimeType !== item2.mimeType) return false;
+    } else if (item1.type === 'text' && item2.type === 'text') {
+      if (item1.text !== item2.text) return false;
+    } else if (item1.type === 'tool-result' && item2.type === 'tool-result') {
+      if (item1.toolCallId !== item2.toolCallId) return false;
+    } else {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+const textToDeltas = (text: string): LanguageModelV1StreamPart[] => {
+  const deltas = text
+    .split(' ')
+    .map((char) => ({ type: 'text-delta' as const, textDelta: `${char} ` }));
+
+  return deltas;
+};
+
+const reasoningToDeltas = (text: string): LanguageModelV1StreamPart[] => {
+  const deltas = text
+    .split(' ')
+    .map((char) => ({ type: 'reasoning' as const, textDelta: `${char} ` }));
+
+  return deltas;
+};
+
+export const getResponseChunksByPrompt = (
+  prompt: CoreMessage[],
+  isReasoningEnabled: boolean = false,
+): Array<LanguageModelV1StreamPart> => {
+  const recentMessage = prompt.at(-1);
+
+  if (!recentMessage) {
+    throw new Error('No recent message found!');
+  }
+
+  if (isReasoningEnabled) {
+    if (compareMessages(recentMessage, TEST_PROMPTS.USER_SKY)) {
+      return [
+        ...reasoningToDeltas('The sky is blue because of rayleigh scattering!'),
+        ...textToDeltas("It's just blue duh!"),
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3 },
+        },
+      ];
+    } else if (compareMessages(recentMessage, TEST_PROMPTS.USER_GRASS)) {
+      return [
+        ...reasoningToDeltas(
+          'Grass is green because of chlorophyll absorption!',
+        ),
+        ...textToDeltas("It's just green duh!"),
+        {
+          type: 'finish',
+          finishReason: 'stop',
+          logprobs: undefined,
+          usage: { completionTokens: 10, promptTokens: 3 },
+        },
+      ];
+    }
+  }
+
+  if (compareMessages(recentMessage, TEST_PROMPTS.USER_THANKS)) {
+    return [
+      ...textToDeltas("You're welcome!"),
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (compareMessages(recentMessage, TEST_PROMPTS.USER_GRASS)) {
+    return [
+      ...textToDeltas("It's just green duh!"),
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (compareMessages(recentMessage, TEST_PROMPTS.USER_SKY)) {
+    return [
+      ...textToDeltas("It's just blue duh!"),
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (compareMessages(recentMessage, TEST_PROMPTS.USER_NEXTJS)) {
+    return [
+      ...textToDeltas('With Next.js, you can ship fast!'),
+
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (
+    compareMessages(recentMessage, TEST_PROMPTS.USER_IMAGE_ATTACHMENT)
+  ) {
+    return [
+      ...textToDeltas('This painting is by Monet!'),
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (compareMessages(recentMessage, TEST_PROMPTS.USER_TEXT_ARTIFACT)) {
+    return [
+      {
+        type: 'tool-call',
+        toolCallId: 'call_123',
+        toolName: 'createDocument',
+        toolCallType: 'function',
+        args: JSON.stringify({
+          title: 'Essay about Silicon Valley',
+          kind: 'text',
+        }),
+      },
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (
+    compareMessages(recentMessage, TEST_PROMPTS.CREATE_DOCUMENT_TEXT_CALL)
+  ) {
+    return [
+      ...textToDeltas(`\n
+# Silicon Valley: The Epicenter of Innovation
+
+## Origins and Evolution
+
+Silicon Valley, nestled in the southern part of the San Francisco Bay Area, emerged as a global technology hub in the late 20th century. Its transformation began in the 1950s when Stanford University encouraged its graduates to start their own companies nearby, leading to the formation of pioneering semiconductor firms that gave the region its name.
+
+## The Innovation Ecosystem
+
+What makes Silicon Valley unique is its perfect storm of critical elements: prestigious universities like Stanford and Berkeley, abundant venture capital, a culture that celebrates risk-taking, and a dense network of talented individuals. This ecosystem has consistently nurtured groundbreaking technologies from personal computers to social media platforms to artificial intelligence.
+
+## Challenges and Criticisms
+
+Despite its remarkable success, Silicon Valley faces significant challenges including extreme income inequality, housing affordability crises, and questions about technology's impact on society. Critics argue the region has developed a monoculture that sometimes struggles with diversity and inclusion.
+
+## Future Prospects
+
+As we move forward, Silicon Valley continues to reinvent itself. While some predict its decline due to remote work trends and competition from other tech hubs, the region's adaptability and innovative spirit suggest it will remain influential in shaping our technological future for decades to come.
+`),
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (
+    compareMessages(recentMessage, TEST_PROMPTS.CREATE_DOCUMENT_TEXT_RESULT)
+  ) {
+    return [
+      {
+        type: 'text-delta',
+        textDelta: 'A document was created and is now visible to the user.',
+      },
+      {
+        type: 'finish',
+        finishReason: 'tool-calls',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (compareMessages(recentMessage, TEST_PROMPTS.GET_WEATHER_CALL)) {
+    return [
+      {
+        type: 'tool-call',
+        toolCallId: 'call_456',
+        toolName: 'getWeather',
+        toolCallType: 'function',
+        args: JSON.stringify({ latitude: 37.7749, longitude: -122.4194 }),
+      },
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  } else if (compareMessages(recentMessage, TEST_PROMPTS.GET_WEATHER_RESULT)) {
+    return [
+      ...textToDeltas('The current temperature in San Francisco is 17°C.'),
+      {
+        type: 'finish',
+        finishReason: 'stop',
+        logprobs: undefined,
+        usage: { completionTokens: 10, promptTokens: 3 },
+      },
+    ];
+  }
+
+  return [{ type: 'text-delta', textDelta: 'Unknown test prompt!' }];
+};
diff --git a/tests/reasoning.test.ts b/tests/reasoning.test.ts
index a70471c..9000285 100644
--- a/tests/reasoning.test.ts
+++ b/tests/reasoning.test.ts
@@ -1,4 +1,3 @@
-import { useAssistant } from '@ai-sdk/react';
 import { ChatPage } from './pages/chat';
 import { test, expect } from '@playwright/test';
 
@@ -11,20 +10,19 @@ test.describe('chat activity with reasoning', () => {
   });
 
   test('send user message and generate response with reasoning', async () => {
-    await chatPage.sendUserMessage('why is the sky blue?');
+    await chatPage.sendUserMessage('Why is the sky blue?');
     await chatPage.waitForMessageGeneration();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
-
-    expect(assistantMessage.content).toBe("it's just blue duh!");
+    expect(assistantMessage.content).toBe("It's just blue duh!");
 
     expect(assistantMessage.reasoning).toBe(
-      'the sky is blue because of rayleigh scattering!',
+      'The sky is blue because of rayleigh scattering!',
     );
   });
 
   test('toggle reasoning visibility', async () => {
-    await chatPage.sendUserMessage('why is the sky blue?');
+    await chatPage.sendUserMessage('Why is the sky blue?');
     await chatPage.waitForMessageGeneration();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
@@ -40,7 +38,7 @@ test.describe('chat activity with reasoning', () => {
   });
 
   test('edit message and resubmit', async () => {
-    await chatPage.sendUserMessage('why is the sky blue?');
+    await chatPage.sendUserMessage('Why is the sky blue?');
     await chatPage.waitForMessageGeneration();
 
     const assistantMessage = await chatPage.getRecentAssistantMessage();
@@ -50,15 +48,15 @@ test.describe('chat activity with reasoning', () => {
 
     const userMessage = await chatPage.getRecentUserMessage();
 
-    await userMessage.edit('why is grass green?');
+    await userMessage.edit('Why is grass green?');
     await chatPage.waitForMessageGeneration();
 
     const updatedAssistantMessage = await chatPage.getRecentAssistantMessage();
 
-    expect(updatedAssistantMessage.content).toBe("it's just green duh!");
+    expect(updatedAssistantMessage.content).toBe("It's just green duh!");
 
     expect(updatedAssistantMessage.reasoning).toBe(
-      'grass is green because of chlorophyll absorption!',
+      'Grass is green because of chlorophyll absorption!',
     );
   });
 });