Replace AI Chat generated content in e2e (#3833)

Author: zenoachtig

packages/gitbook/e2e/internal.spec.ts

@@ -39,6 +39,29 @@ import {
 
 const AI_PROMPT = `You're being invoked by the GitBook CI/CD pipeline. To make screenshot testing of the GitBook Assistant visually consistent, look up the title of the first page you find and respond with only EXACTLY its title. To find the page title, invoke the search tool with the query "GitBook". Before invoking the search tool, respond with the exact text: "I'm going to look up 'GitBook' and then respond with only the page title.". Do not execute any other tools or output any other text.`;
 
+const overrideAIInitialState = () => {
+    const greeting = document.querySelector('[data-testid="ai-chat-time-greeting"]');
+    if (greeting) {
+        greeting.textContent = 'Good morning';
+    }
+};
+const overrideAIResponse = () => {
+    const userMessage = document.querySelector('[data-testid="ai-chat-message-user"]');
+    if (userMessage) {
+        userMessage.textContent = '[Replaced message] Chat message sent by the user';
+    }
+    const assistantMessage = document.querySelectorAll(
+        '[data-testid="ai-chat-message-assistant"] .ai-response-document'
+    );
+    assistantMessage.forEach((message) => {
+        message.innerHTML = '[Replaced message] AI chat response';
+    });
+    const suggestions = document.querySelectorAll('[data-testid="ai-chat-followup-suggestion"]');
+    suggestions.forEach((suggestion) => {
+        suggestion.textContent = 'Follow-up suggestion';
+    });
+};
+
 const searchTestCases: Test[] = [
     {
         name: 'Search - AI Mode: None - Complete flow',
@@ -147,14 +170,7 @@ const searchTestCases: Test[] = [
                 timeout: 60_000,
             });
             // Override text content for visual consistency in screenshots
-            await page.evaluate(() => {
-                const suggestions = document.querySelectorAll(
-                    '[data-testid="ai-chat-followup-suggestion"]'
-                );
-                suggestions.forEach((suggestion) => {
-                    suggestion.textContent = 'Follow-up suggestion';
-                });
-            });
+            await page.evaluate(overrideAIResponse);
         },
     },
     {
@@ -169,12 +185,7 @@ const searchTestCases: Test[] = [
             await expect(page.getByTestId('ai-chat')).toBeVisible();
             await expect(page.getByTestId('ai-chat-input')).toBeFocused();
             // Override text content for visual consistency in screenshots
-            await page.evaluate(() => {
-                const greeting = document.querySelector('[data-testid="ai-chat-time-greeting"]');
-                if (greeting) {
-                    greeting.textContent = 'Good morning';
-                }
-            });
+            await page.evaluate(overrideAIInitialState);
         },
     },
     {
@@ -190,12 +201,7 @@ const searchTestCases: Test[] = [
             await expect(page.getByTestId('ai-chat')).toBeVisible();
             await expect(page.getByTestId('ai-chat-input')).toBeFocused();
             // Override text content for visual consistency in screenshots
-            await page.evaluate(() => {
-                const greeting = document.querySelector('[data-testid="ai-chat-time-greeting"]');
-                if (greeting) {
-                    greeting.textContent = 'Good morning';
-                }
-            });
+            await page.evaluate(overrideAIInitialState);
         },
     },
     {
@@ -211,12 +217,7 @@ const searchTestCases: Test[] = [
             await expect(page.getByTestId('ai-chat')).toBeVisible();
             await expect(page.getByTestId('ai-chat-input')).toBeFocused();
             // Override text content for visual consistency in screenshots
-            await page.evaluate(() => {
-                const greeting = document.querySelector('[data-testid="ai-chat-time-greeting"]');
-                if (greeting) {
-                    greeting.textContent = 'Good morning';
-                }
-            });
+            await page.evaluate(overrideAIInitialState);
         },
     },
     {
@@ -236,14 +237,7 @@ const searchTestCases: Test[] = [
                 timeout: 60_000,
             });
             // Override text content for visual consistency in screenshots
-            await page.evaluate(() => {
-                const suggestions = document.querySelectorAll(
-                    '[data-testid="ai-chat-followup-suggestion"]'
-                );
-                suggestions.forEach((suggestion) => {
-                    suggestion.textContent = 'Follow-up suggestion';
-                });
-            });
+            await page.evaluate(overrideAIResponse);
         },
     },
 ];

packages/gitbook/src/components/AI/server-actions/AIMessageView.tsx

@@ -35,7 +35,7 @@ export function AIMessageView(
                                 wrapBlocksInSuspense: false,
                                 withLinkPreviews,
                             }}
-                            style="mt-2 space-y-4 *:origin-top-left *:animate-blur-in-slow"
+                            style="ai-response-document mt-2 space-y-4 *:origin-top-left *:animate-blur-in-slow"
                         />
 
                         {withToolCalls && step.toolCalls && step.toolCalls.length > 0 ? (