diff --git a/packages/matrix/helpers/index.ts b/packages/matrix/helpers/index.ts
index 2a42716be1..8684f848b9 100644
--- a/packages/matrix/helpers/index.ts
+++ b/packages/matrix/helpers/index.ts
@@ -73,13 +73,8 @@ export async function reloadAndOpenAiAssistant(page: Page) {
 
 export async function openAiAssistant(page: Page) {
   await page.locator('[data-test-open-ai-assistant]').click();
-  await page.waitForFunction(() =>
-    document.querySelector('[data-test-close-ai-assistant]'),
-  );
-
-  await page.waitForFunction(() =>
-    document.querySelector('[data-test-room]')?.getAttribute('data-test-room'),
-  ); // Opening the AI assistant either opens last room or creates one - wait for it to settle
+  await expect(page.locator('[data-test-close-ai-assistant]')).toHaveCount(1);
+  await expect(page.locator('[data-test-room]')).toHaveCount(1);
 }
 
 export async function createRealm(
@@ -265,9 +260,7 @@ export async function login(
 ) {
   await openRoot(page, opts?.url);
 
-  await page.waitForFunction(() =>
-    document.querySelector('[data-test-username-field]'),
-  );
+  await expect(page.locator('[data-test-username-field]')).toBeEditable();
   await page.locator('[data-test-username-field]').fill(username);
   await page.locator('[data-test-password-field]').fill(password);
   await page.locator('[data-test-login-btn]').click();
diff --git a/packages/matrix/tests/room-creation.spec.ts b/packages/matrix/tests/room-creation.spec.ts
index 4e080c2988..ba89fa2dac 100644
--- a/packages/matrix/tests/room-creation.spec.ts
+++ b/packages/matrix/tests/room-creation.spec.ts
@@ -235,10 +235,7 @@ test.describe('Room creation', () => {
 
     // Open assistant without waiting for [data-test-room] which won’t show on a new account
     await page.locator('[data-test-open-ai-assistant]').click();
-    await page.waitForFunction(() =>
-      document.querySelector('[data-test-close-ai-assistant]'),
-    );
-
+    await expect(page.locator(`[data-test-close-ai-assistant]`)).toHaveCount(1);
     await expect(page.locator(`[data-test-chat-title]`)).not.toHaveText(
       'New AI Assistant Chat',
     );