diff --git a/tests/blocksuite/e2e/utils/actions/edgeless.ts b/tests/blocksuite/e2e/utils/actions/edgeless.ts
index f00007bd76..0c3ab90e65 100644
--- a/tests/blocksuite/e2e/utils/actions/edgeless.ts
+++ b/tests/blocksuite/e2e/utils/actions/edgeless.ts
@@ -376,28 +376,19 @@ export async function setEdgelessTool(
         'shape',
         false
       );
-      const shapeToolBox = await shapeToolButton.boundingBox();
-      if (!shapeToolBox) {
-        throw new Error('shapeToolBox is not found');
-      }
-
-      await page.mouse.click(shapeToolBox.x + 2, shapeToolBox.y + 2);
+      await shapeToolButton.click({ position: { x: 2, y: 2 } });
 
       const shapeMenu = page.locator('edgeless-shape-menu');
+      if ((await shapeMenu.count()) === 0 && shape === Shape.Square) {
+        break;
+      }
       await expect(shapeMenu).toBeVisible();
 
       const squareShapeButton = shapeMenu
         .locator('edgeless-tool-icon-button')
         .filter({ hasText: shape });
       await expect(squareShapeButton).toBeVisible();
-      const squareShapeBox = await squareShapeButton.boundingBox();
-      if (!squareShapeBox) {
-        throw new Error('squareShapeBox is not found');
-      }
-      await page.mouse.click(
-        squareShapeBox.x + squareShapeBox.width / 2,
-        squareShapeBox.y + squareShapeBox.height / 2
-      );
+      await squareShapeButton.click();
       break;
     }
   }
diff --git a/tests/blocksuite/e2e/utils/actions/misc.ts b/tests/blocksuite/e2e/utils/actions/misc.ts
index 6e6b7f96a1..bfcff6b6e7 100644
--- a/tests/blocksuite/e2e/utils/actions/misc.ts
+++ b/tests/blocksuite/e2e/utils/actions/misc.ts
@@ -552,6 +552,27 @@ export async function focusRichText(
   await expect(locator).toBeVisible();
   // need to set `force` to true when clicking on `affine-selected-blocks`
   await locator.click({ force: true, position: options?.clickPosition });
+  await expect
+    .poll(() =>
+      page.evaluate(
+        ([i, currentEditorIndex]) => {
+          const editorHost =
+            document.querySelectorAll('editor-host')[currentEditorIndex];
+          const richText =
+            editorHost?.querySelectorAll<RichText>('rich-text')[i];
+          const inlineEditor = richText?.inlineEditor;
+          if (!inlineEditor) {
+            return false;
+          }
+          if (inlineEditor.getInlineRange() === null) {
+            inlineEditor.focusIndex(0);
+          }
+          return inlineEditor.getInlineRange() !== null;
+        },
+        [i, currentEditorIndex]
+      )
+    )
+    .toBe(true);
 }
 
 export async function focusRichTextEnd(page: Page, i = 0) {