diff --git a/tests/affine-cloud-copilot/e2e/ai-action/generate-an-image-with-text.spec.ts b/tests/affine-cloud-copilot/e2e/ai-action/generate-an-image-with-text.spec.ts
index eb5b3e488..eec01e1fd 100644
--- a/tests/affine-cloud-copilot/e2e/ai-action/generate-an-image-with-text.spec.ts
+++ b/tests/affine-cloud-copilot/e2e/ai-action/generate-an-image-with-text.spec.ts
@@ -51,6 +51,21 @@ test.describe('AIAction/GenerateAnImageWithText', () => {
     expect(responses).toEqual(new Set(['insert-below']));
   });
 
+  test('should generate an image for the selected shape in edgeless', async ({
+    page,
+    utils,
+  }) => {
+    const { generateImage } = await utils.editor.askAIWithEdgeless(
+      page,
+      async () => {
+        await utils.editor.createShape(page, 'HelloWorld');
+      }
+    );
+    const { answer, responses } = await generateImage();
+    await expect(answer.getByTestId('ai-answer-image')).toBeVisible();
+    expect(responses).toEqual(new Set(['insert-below']));
+  });
+
   test('should show chat history in chat panel', async ({ page, utils }) => {
     const { generateImage } = await utils.editor.askAIWithText(page, 'Panda');
     const { answer } = await generateImage();
diff --git a/tests/affine-cloud-copilot/e2e/ai-action/generate-headings.spec.ts b/tests/affine-cloud-copilot/e2e/ai-action/generate-headings.spec.ts
index b6d1a1906..1b2f9ec8c 100644
--- a/tests/affine-cloud-copilot/e2e/ai-action/generate-headings.spec.ts
+++ b/tests/affine-cloud-copilot/e2e/ai-action/generate-headings.spec.ts
@@ -29,53 +29,53 @@ test.describe('AIAction/GenerateHeadings', () => {
     expect(responses).toEqual(new Set(['insert-above', 'replace-selection']));
   });
 
-  test.fixme(
-    'should generate headings for selected text block in edgeless',
-    async ({ page, utils }) => {
-      const { generateHeadings } = await utils.editor.askAIWithEdgeless(
-        page,
-        async () => {
-          await utils.editor.createEdgelessText(
-            page,
-            'AFFiNE is a workspace with fully merged docs'
-          );
-        }
-      );
+  test('should generate headings for selected text block in edgeless', async ({
+    page,
+    utils,
+  }) => {
+    const { generateHeadings } = await utils.editor.askAIWithEdgeless(
+      page,
+      async () => {
+        await utils.editor.createEdgelessText(
+          page,
+          'AFFiNE is a workspace with fully merged docs'
+        );
+      }
+    );
 
-      const { answer, responses } = await generateHeadings();
-      await Promise.race([
-        answer.locator('h1').isVisible(),
-        answer.locator('h2').isVisible(),
-        answer.locator('h3').isVisible(),
-      ]);
-      await expect(answer).toHaveText(/AFFiNE/, { timeout: 10000 });
-      expect(responses).toEqual(new Set(['insert-above']));
-    }
-  );
+    const { answer, responses } = await generateHeadings();
+    await Promise.race([
+      answer.locator('h1').isVisible(),
+      answer.locator('h2').isVisible(),
+      answer.locator('h3').isVisible(),
+    ]);
+    await expect(answer).toHaveText(/AFFiNE/, { timeout: 10000 });
+    expect(responses).toEqual(new Set(['insert-below']));
+  });
 
-  test.fixme(
-    'should generate headings for selected note block in edgeless',
-    async ({ page, utils }) => {
-      const { generateHeadings } = await utils.editor.askAIWithEdgeless(
-        page,
-        async () => {
-          await utils.editor.createEdgelessNote(
-            page,
-            'AFFiNE is a workspace with fully merged docs'
-          );
-        }
-      );
+  test('should generate headings for selected note block in edgeless', async ({
+    page,
+    utils,
+  }) => {
+    const { generateHeadings } = await utils.editor.askAIWithEdgeless(
+      page,
+      async () => {
+        await utils.editor.createEdgelessNote(
+          page,
+          'AFFiNE is a workspace with fully merged docs'
+        );
+      }
+    );
 
-      const { answer, responses } = await generateHeadings();
-      await Promise.race([
-        answer.locator('h1').isVisible(),
-        answer.locator('h2').isVisible(),
-        answer.locator('h3').isVisible(),
-      ]);
-      await expect(answer).toHaveText(/AFFiNE/, { timeout: 10000 });
-      expect(responses).toEqual(new Set(['insert-above']));
-    }
-  );
+    const { answer, responses } = await generateHeadings();
+    await Promise.race([
+      answer.locator('h1').isVisible(),
+      answer.locator('h2').isVisible(),
+      answer.locator('h3').isVisible(),
+    ]);
+    await expect(answer).toHaveText(/AFFiNE/, { timeout: 10000 });
+    expect(responses).toEqual(new Set(['insert-below']));
+  });
 
   test('should show chat history in chat panel', async ({ page, utils }) => {
     const { generateHeadings } = await utils.editor.askAIWithText(
diff --git a/tests/affine-cloud-copilot/e2e/ai-action/make-it-real.spec.ts b/tests/affine-cloud-copilot/e2e/ai-action/make-it-real.spec.ts
index 86d33201d..ea8078d4e 100644
--- a/tests/affine-cloud-copilot/e2e/ai-action/make-it-real.spec.ts
+++ b/tests/affine-cloud-copilot/e2e/ai-action/make-it-real.spec.ts
@@ -60,6 +60,22 @@ test.describe('AIAction/MakeItReal', () => {
     expect(responses).toEqual(new Set(['insert-below']));
   });
 
+  test('should support making the selected element to real in edgeless', async ({
+    page,
+    utils,
+  }) => {
+    const { makeItReal } = await utils.editor.askAIWithEdgeless(
+      page,
+      async () => {
+        await utils.editor.createShape(page, 'HelloWorld');
+      }
+    );
+
+    const { answer, responses } = await makeItReal();
+    await expect(answer.locator('iframe')).toBeVisible({ timeout: 30000 });
+    expect(responses).toEqual(new Set(['insert-below']));
+  });
+
   test('should show chat history in chat panel', async ({ page, utils }) => {
     const { makeItReal } = await utils.editor.askAIWithText(
       page,
diff --git a/tests/affine-cloud-copilot/e2e/chat-with/edgeless-shape.spec.ts b/tests/affine-cloud-copilot/e2e/chat-with/edgeless-shape.spec.ts
deleted file mode 100644
index b167aa5d8..000000000
--- a/tests/affine-cloud-copilot/e2e/chat-with/edgeless-shape.spec.ts
+++ /dev/null
@@ -1,3 +0,0 @@
-import { test } from '../base/base-test';
-
-test.describe('AIChatWith/EdgelessShape', () => {});
diff --git a/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts b/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts
index 3d46057e8..7ace8b9c2 100644
--- a/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts
+++ b/tests/affine-cloud-copilot/e2e/utils/editor-utils.ts
@@ -160,6 +160,33 @@ export class EditorUtils {
     return id;
   }
 
+  public static async createShape(page: Page, text: string) {
+    // Create shape
+    await page.keyboard.press('s');
+    await page.mouse.click(400, 400);
+
+    const id = await page.evaluate(() => {
+      const edgelessBlock = document.querySelector(
+        'affine-edgeless-root'
+      ) as EdgelessRootBlockComponent;
+      if (!edgelessBlock) {
+        throw new Error('edgeless block not found');
+      }
+      const shapes = edgelessBlock.gfx.gfxElements.filter(
+        (el: GfxModel) => 'type' in el && el.type === 'shape'
+      );
+
+      return shapes[shapes.length - 1].id;
+    });
+
+    // Insert text inside shape
+    await page.mouse.dblclick(400, 400);
+    await page.keyboard.insertText(text);
+    // Prevent the shape from being dragged
+    await page.mouse.click(500, 500);
+    return id;
+  }
+
   public static async getMindMapNode(
     page: Page,
     mindmapId: string,