From 0325922547082d7b353891b6dbf40876b9f03a57 Mon Sep 17 00:00:00 2001
From: Mario Ruci <mario.ruci@dfinity.org>
Date: Fri, 18 Sep 2026 19:07:54 +0200
Subject: [PATCH] test(e2e): bound the waits in the disaster recovery test

The test shares one budget across a long flow and contained two loops that
could not fail on their own: one polling the station module hash, one polling
for the restored balance. Whichever step happened to be slowest consumed the
budget, and the failure was reported as a timeout inside whichever wait the
test was sitting in, which is why the reported line moves between steps
across runs.

Both loops become bounded polls with their own timeouts and messages, so a
step that does not complete fails against that step and reports the value it
last read. The entry point into the recovery page is also waited for rather
than clicked immediately, since it only appears once the wallet has noticed
the station is uninstalled.
---
 tests/e2e/disaster-recovery.spec.ts     | 34 +++++++++++++++----------
 tests/e2e/page-objects/settings.page.ts | 18 ++++++++-----
 2 files changed, 32 insertions(+), 20 deletions(-)

diff --git a/tests/e2e/disaster-recovery.spec.ts b/tests/e2e/disaster-recovery.spec.ts
index 154928dad..86edd6817 100644
--- a/tests/e2e/disaster-recovery.spec.ts
+++ b/tests/e2e/disaster-recovery.spec.ts
@@ -49,7 +49,12 @@ test('can recover uninstalled station', async ({ page }) => {
   );
 
   await page.goto(walletUrl);
-  await page.getByText(/disaster recovery/i).click();
+
+  // The entry point only appears once the wallet has noticed the station is uninstalled, so it is
+  // waited for rather than clicked straight away.
+  const disasterRecoveryLink = page.getByText(/disaster recovery/i).first();
+  await disasterRecoveryLink.waitFor({ state: 'visible', timeout: 60_000 });
+  await disasterRecoveryLink.click();
 
   const disasterRecoveryPage = new DisasterRecoveryPage(page);
   await disasterRecoveryPage.assertIsOn();
@@ -64,17 +69,20 @@ test('can recover uninstalled station', async ({ page }) => {
   await accountsPage.openByName('Main');
 
   await accountPage.pickByAsset('ICP');
-  await accountAssetPage.getBalance();
-
-  while (true) {
-    // refresh the page
-    await page.reload();
-    const balance = await accountAssetPage.getBalance();
-
-    if (balance!.includes('5.0')) {
-      break;
-    }
 
-    await page.waitForTimeout(5000);
-  }
+  // Bounded so that a balance that never returns fails here reporting what it actually read,
+  // instead of spinning until the shared test budget runs out somewhere unrelated.
+  await expect
+    .poll(
+      async () => {
+        await page.reload();
+        return (await accountAssetPage.getBalance()) ?? '';
+      },
+      {
+        message: 'the ICP balance should be restored after disaster recovery',
+        timeout: 180_000,
+        intervals: [5_000],
+      },
+    )
+    .toContain('5.0');
 });
diff --git a/tests/e2e/page-objects/settings.page.ts b/tests/e2e/page-objects/settings.page.ts
index 13cc0d96a..00b46d5af 100644
--- a/tests/e2e/page-objects/settings.page.ts
+++ b/tests/e2e/page-objects/settings.page.ts
@@ -1,4 +1,4 @@
-import { Page } from '@playwright/test';
+import { expect, Page } from '@playwright/test';
 import { getWalletPath } from '../config';
 import { getCanisterInfo } from '../utils/dfx.utils';
 
@@ -25,12 +25,16 @@ export class SettingsPage {
     await this.page.getByTestId('continue-action-btn').click();
     await this.page.getByTestId('submit-action-btn').click();
 
-    while (checkForNewModuleHash) {
-      const newModuleHash = getCanisterInfo(stationId).moduleHash;
-      if (newModuleHash !== originalModuleHash) {
-        break;
-      }
-      await this.page.waitForTimeout(1000);
+    if (checkForNewModuleHash) {
+      // Bounded so a wasm that never installs fails here rather than consuming the whole test
+      // budget and surfacing as a timeout in a later, unrelated step.
+      await expect
+        .poll(() => getCanisterInfo(stationId).moduleHash, {
+          message: 'the station module hash should change after installing the custom wasm',
+          timeout: 120_000,
+          intervals: [2_000],
+        })
+        .not.toBe(originalModuleHash);
     }
 
     // wait till the canister starts again