test(core): chat with collection & tags (#11388)

### TL:DR

* AI chat with collection E2E
* AI chat with tag E2E

> Close BS-3007

packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts

@@ -607,7 +607,13 @@ export class ChatPanel extends SignalWatcher(
 
     return html`<div class="chat-panel-container" style=${style}>
       <div class="chat-panel-title">
-        <div>${isEmbedding ? `Embedding ${done}/${total}` : 'AFFiNE AI'}</div>
+        <div>
+          ${isEmbedding
+            ? html`<span data-testid="chat-panel-embedding-progress"
+                >Embedding ${done}/${total}</span
+              >`
+            : 'AFFiNE AI'}
+        </div>
         <div
           @click=${() => {
             AIProvider.toggleGeneralAIOnboarding?.(true);

packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/add-popover.ts

@@ -254,6 +254,9 @@ export class ChatPanelAddPopover extends SignalWatcher(
   @property({ attribute: false })
   accessor abortController!: AbortController;
 
+  @property({ attribute: 'data-testid', reflect: true })
+  accessor testId: string = 'ai-search-input';
+
   @query('.search-input')
   accessor searchInput!: HTMLInputElement;
 

packages/frontend/core/src/modules/explorer/views/nodes/collection/operations.tsx

@@ -118,6 +118,7 @@ export const useExplorerCollectionNodeOperations = (
         view: (
           <IconButton
             size="16"
+            data-testid="collection-add-doc-button"
             onClick={handleAddDocToCollection}
             tooltip={t[
               'com.affine.rootAppSidebar.explorer.collection-add-tooltip'

packages/frontend/core/src/modules/explorer/views/nodes/tag/operations.tsx

@@ -91,6 +91,7 @@ export const useExplorerTagNodeOperations = (
           <IconButton
             size="16"
             onClick={handleNewDoc}
+            data-testid="tag-add-doc-button"
             tooltip={t['com.affine.rootAppSidebar.explorer.tag-add-tooltip']()}
           >
             <PlusIcon />

packages/frontend/core/src/modules/explorer/views/sections/tags/index.tsx

@@ -43,6 +43,7 @@ export const ExplorerTags = () => {
   return (
     <CollapsibleSection
       name="tags"
+      testId="explorer-tags"
       headerClassName={styles.draggedOverHighlight}
       title={t['com.affine.rootAppSidebar.tags']()}
       actions={

tests/affine-cloud-copilot/e2e/basic/authority.spec.ts

@@ -4,7 +4,7 @@ import { test } from '../base/base-test';
 
 test.describe('AIBasic/Authority', () => {
   test.beforeEach(async ({ page, utils }) => {
-    await utils.testUtils.setupTestEnvironment(page);
+    await utils.testUtils.setupTestEnvironment(page, false);
     await utils.chatPanel.openChatPanel(page);
   });
 

tests/affine-cloud-copilot/e2e/basic/chat.spec.ts

@@ -256,7 +256,6 @@ test.describe('AIBasic/Chat', () => {
     ]);
 
     const { actions } = await utils.chatPanel.getLatestAssistantMessage(page);
-    await page.pause();
     await actions.retry();
 
     await utils.chatPanel.waitForHistory(page, [

tests/affine-cloud-copilot/e2e/chat-with/collections.spec.ts

@@ -1,3 +1,79 @@
+import { loginUser } from '@affine-test/kit/utils/cloud';
+import { expect } from '@playwright/test';
+
 import { test } from '../base/base-test';
 
-test.describe('AIChatWith/Collections', () => {});
+test.describe('AIChatWith/Collections', () => {
+  test.beforeEach(async ({ page, utils }) => {
+    const user = await utils.testUtils.getUser();
+    await loginUser(page, user);
+    await utils.testUtils.setupTestEnvironment(page);
+    await utils.chatPanel.openChatPanel(page);
+
+    // Create two collections
+    await utils.editor.createCollectionAndDoc(
+      page,
+      'Collection 1',
+      'EEee is a cute cat'
+    );
+    await utils.editor.createCollectionAndDoc(
+      page,
+      'Collection 2',
+      'FFff is a cute dog'
+    );
+  });
+
+  test('should support chat with collection', async ({ page, utils }) => {
+    await utils.chatPanel.chatWithCollections(page, ['Collection 1']);
+    await utils.chatPanel.makeChat(page, 'What is EEee(Use English)');
+    await utils.chatPanel.waitForHistory(page, [
+      {
+        role: 'user',
+        content: 'What is EEee(Use English)',
+      },
+      {
+        role: 'assistant',
+        status: 'success',
+      },
+    ]);
+
+    await expect(async () => {
+      const { content, message } =
+        await utils.chatPanel.getLatestAssistantMessage(page);
+      expect(content).toMatch(/EEee.*cat/);
+      expect(await message.locator('affine-footnote-node').count()).toBe(1);
+    }).toPass();
+  });
+
+  test('should support chat with multiple collections', async ({
+    page,
+    utils,
+  }) => {
+    await utils.chatPanel.chatWithCollections(page, [
+      'Collection 1',
+      'Collection 2',
+    ]);
+    await utils.chatPanel.makeChat(
+      page,
+      'What is EEee? What is FFff?(Use English)'
+    );
+    await utils.chatPanel.waitForHistory(page, [
+      {
+        role: 'user',
+        content: 'What is EEee? What is FFff?(Use English)',
+      },
+      {
+        role: 'assistant',
+        status: 'success',
+      },
+    ]);
+
+    await expect(async () => {
+      const { content, message } =
+        await utils.chatPanel.getLatestAssistantMessage(page);
+      expect(content).toMatch(/EEee.*cat/);
+      expect(content).toMatch(/FFff.*dog/);
+      expect(await message.locator('affine-footnote-node').count()).toBe(2);
+    }).toPass();
+  });
+});

tests/affine-cloud-copilot/e2e/chat-with/tags.spec.ts

@@ -1,3 +1,61 @@
+import { loginUser } from '@affine-test/kit/utils/cloud';
+import { expect } from '@playwright/test';
+
 import { test } from '../base/base-test';
 
-test.describe('AIChatWith/tags', () => {});
+test.describe('AIChatWith/tags', () => {
+  test.beforeEach(async ({ page, utils }) => {
+    const user = await utils.testUtils.getUser();
+    await loginUser(page, user);
+    await utils.testUtils.setupTestEnvironment(page);
+    await utils.chatPanel.openChatPanel(page);
+    await utils.editor.createTagAndDoc(page, 'Tag 1', 'EEee is a cute cat');
+    await utils.editor.createTagAndDoc(page, 'Tag 2', 'FFff is a cute dog');
+  });
+
+  test('should support chat with tag', async ({ page, utils }) => {
+    await utils.chatPanel.chatWithTags(page, ['Tag 1']);
+    await utils.chatPanel.makeChat(page, 'What is EEee(Use English)');
+    await utils.chatPanel.waitForHistory(page, [
+      {
+        role: 'user',
+        content: 'What is EEee(Use English)',
+      },
+      {
+        role: 'assistant',
+        status: 'success',
+      },
+    ]);
+    await expect(async () => {
+      const { content, message } =
+        await utils.chatPanel.getLatestAssistantMessage(page);
+      expect(content).toMatch(/EEee.*cat/);
+      expect(await message.locator('affine-footnote-node').count()).toBe(1);
+    }).toPass();
+  });
+
+  test('should support chat with multiple tags', async ({ page, utils }) => {
+    await utils.chatPanel.chatWithTags(page, ['Tag 1', 'Tag 2']);
+    await utils.chatPanel.makeChat(
+      page,
+      'What is EEee? What is FFff?(Use English)'
+    );
+    await utils.chatPanel.waitForHistory(page, [
+      {
+        role: 'user',
+        content: 'What is EEee? What is FFff?(Use English)',
+      },
+      {
+        role: 'assistant',
+        status: 'success',
+      },
+    ]);
+    await expect(async () => {
+      const { content, message } =
+        await utils.chatPanel.getLatestAssistantMessage(page);
+      expect(content).toMatch(/EEee.*cat/);
+      expect(content).toMatch(/FFff.*dog/);
+      expect(await message.locator('affine-footnote-node').count()).toBe(2);
+    }).toPass();
+  });
+});

tests/affine-cloud-copilot/e2e/utils/chat-panel-utils.ts

@@ -225,7 +225,6 @@ export class ChatPanelUtils {
       await expect(states.every(state => state === 'finished')).toBe(true);
     }).toPass({ timeout: 20000 });
 
-    await page.pause();
     await this.makeChat(page, text);
   }
 
@@ -249,6 +248,39 @@ export class ChatPanelUtils {
     await this.makeChat(page, text);
   }
 
+  public static async chatWithTags(page: Page, tags: string[]) {
+    for (const tag of tags) {
+      const withButton = await page.getByTestId('chat-panel-with-button');
+      await withButton.click();
+      const withMenu = await page.getByTestId('ai-add-popover');
+      await withMenu.getByTestId('ai-chat-with-tags').click();
+      await withMenu.getByText(tag).click();
+      await page.getByTestId('chat-panel-chips').getByText(tag);
+    }
+    await this.waitForEmbeddingProgress(page);
+  }
+
+  public static async chatWithCollections(page: Page, collections: string[]) {
+    for (const collection of collections) {
+      const withButton = await page.getByTestId('chat-panel-with-button');
+      await withButton.click();
+      const withMenu = await page.getByTestId('ai-add-popover');
+      await withMenu.getByTestId('ai-chat-with-collections').click();
+      await withMenu.getByText(collection).click();
+      await page.getByTestId('chat-panel-chips').getByText(collection);
+    }
+    await this.waitForEmbeddingProgress(page);
+  }
+
+  public static async waitForEmbeddingProgress(page: Page) {
+    await page.getByTestId('chat-panel-embedding-progress').waitFor({
+      state: 'visible',
+    });
+    await page.getByTestId('chat-panel-embedding-progress').waitFor({
+      state: 'hidden',
+    });
+  }
+
   public static async enableNetworkSearch(page: Page) {
     const networkSearch = await page.getByTestId('chat-network-search');
     if ((await networkSearch.getAttribute('data-active')) === 'false') {

tests/affine-cloud-copilot/e2e/utils/editor-utils.ts

@@ -209,6 +209,63 @@ export class EditorUtils {
     );
   }
 
+  public static async createCollectionAndDoc(
+    page: Page,
+    collectionName: string,
+    docContent: string
+  ) {
+    // Create collection
+    await page.getByTestId('explorer-bar-add-collection-button').click();
+    const input = await page.getByTestId('prompt-modal-input');
+    await input.focus();
+    await input.pressSequentially(collectionName);
+    await page.getByTestId('prompt-modal-confirm').click();
+    const collections = await page.getByTestId('collapsible-section-content');
+    const collection = await collections
+      .getByText(collectionName)
+      .locator('..');
+
+    // Create doc
+    await collection.hover();
+    await collection.getByTestId('collection-add-doc-button').click();
+    await page.getByTestId('confirm-modal-confirm').click();
+    await this.focusToEditor(page);
+    const texts = docContent.split('\n');
+    for (const [index, line] of texts.entries()) {
+      await page.keyboard.insertText(line);
+      if (index !== texts.length - 1) {
+        await page.keyboard.press('Enter');
+      }
+    }
+  }
+
+  public static async createTagAndDoc(
+    page: Page,
+    tagName: string,
+    docContent: string
+  ) {
+    // Create tag
+    const tags = await page.getByTestId('explorer-tags');
+    await tags.getByTestId('explorer-bar-add-favorite-button').click();
+    const input = await page.getByTestId('rename-modal-input');
+    await input.focus();
+    await input.pressSequentially(tagName);
+    await input.press('Enter');
+    const tag = await tags.getByText(tagName).locator('..');
+
+    // Create doc
+    await tag.hover();
+    await tag.getByTestId('tag-add-doc-button').click();
+    await this.focusToEditor(page);
+    const texts = docContent.split('\n');
+    for (const [index, line] of texts.entries()) {
+      await page.keyboard.insertText(line);
+      if (index !== texts.length - 1) {
+        await page.keyboard.press('Enter');
+      }
+    }
+  }
+
   public static async selectElementInEdgeless(page: Page, elements: string[]) {
     await page.evaluate(
       ({ elements }) => {

tests/affine-cloud-copilot/e2e/utils/test-utils.ts

@@ -1,4 +1,7 @@
-import { createRandomAIUser } from '@affine-test/kit/utils/cloud';
+import {
+  createRandomAIUser,
+  enableCloudWorkspace,
+} from '@affine-test/kit/utils/cloud';
 import { openHomePage, setCoreUrl } from '@affine-test/kit/utils/load-page';
 import {
   clickNewPageButton,
@@ -52,10 +55,13 @@ export class TestUtils {
     };
   }
 
-  public async setupTestEnvironment(page: Page) {
+  public async setupTestEnvironment(page: Page, enableCloud: boolean = true) {
     await openHomePage(page);
     await clickNewPageButton(page);
     await waitForEditorLoad(page);
+    if (enableCloud) {
+      await enableCloudWorkspace(page);
+    }
   }
 
   public async createTestWorkspace(page: Page, name: string = 'test') {