From ced11b6be51f7e68af5b8a03f7d780c4c2931d1d Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Thu, 9 Jan 2025 16:03:13 -0500
Subject: [PATCH] tests: fix e2e shutdown test (#3390)

---
 frontend/e2e-tests/shutdown.spec.ts | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/frontend/e2e-tests/shutdown.spec.ts b/frontend/e2e-tests/shutdown.spec.ts
index 7ba3bd79e1c..93b90fa45cc 100644
--- a/frontend/e2e-tests/shutdown.spec.ts
+++ b/frontend/e2e-tests/shutdown.spec.ts
@@ -36,9 +36,20 @@ test("restart kernel", async ({ page }) => {
   const appUrl = getAppUrl("shutdown.py");
   await page.goto(appUrl);
 
+  // Wait for page to be fully loaded
+  await page.waitForLoadState("networkidle");
+
   await page.getByTestId("notebook-menu-dropdown").click();
-  await page.getByText("Restart kernel").click();
-  await page.getByLabel("Confirm Restart").click();
+  // Wait for dropdown to be visible and stable
+  await page.waitForTimeout(100);
+
+  const restartButton = page.getByRole("menuitem", { name: "Restart kernel" });
+  await restartButton.waitFor({ state: "visible" });
+  await restartButton.click();
+
+  const confirmButton = page.getByRole("button", { name: "Confirm Restart" });
+  await confirmButton.waitFor({ state: "visible" });
+  await confirmButton.click();
 
   await expect(page.getByText("None", { exact: true })).toBeVisible();
 });