diff --git a/ui-tests/package.json b/ui-tests/package.json
index 2b51afa..0c6e38f 100644
--- a/ui-tests/package.json
+++ b/ui-tests/package.json
@@ -10,7 +10,7 @@
         "test:debug": "PWDEBUG=1 npx playwright test"
     },
     "devDependencies": {
-        "@jupyterlab/galata": "^5.0.8",
+        "@jupyterlab/galata": "^5.1.0",
         "@playwright/test": "^1.32.0",
         "@types/klaw-sync": "^6.0.1"
     },
diff --git a/ui-tests/tests/ui.spec.ts b/ui-tests/tests/ui.spec.ts
index bcdbc43..c480fec 100644
--- a/ui-tests/tests/ui.spec.ts
+++ b/ui-tests/tests/ui.spec.ts
@@ -119,18 +119,17 @@ test.describe('UI Test', () => {
       await page.notebook.openByPath(fullPath);
       await page.notebook.activate(fullPath);
       await page.locator('div.jpcad-Spinner').waitFor({ state: 'hidden' });
-      const btn = await page.locator(
-        "button.jp-ToolbarButtonComponent[data-command='jupytercad:newBox']"
-      );
-      await btn.click();
+
+      await page.getByTitle('New Box').click();
+
       const nameInput = await page.locator(
         'input[id^="id-jp-schemaform"][label="Name"]'
       );
-      nameInput.fill('Foo');
+      await nameInput.fill('Foo');
       const accept = await page.locator('div.jp-Dialog-buttonLabel', {
         hasText: 'Submit'
       });
-      accept.click();
+      await accept.click();
 
       await page
         .getByRole('tablist', { name: 'main sidebar' })
@@ -248,10 +247,8 @@ test.describe('UI Test', () => {
       await page.locator('div.jpcad-Spinner').waitFor({ state: 'hidden' });
 
       // Create a cone
-      const btn = await page.locator(
-        "button.jp-ToolbarButtonComponent[data-command='jupytercad:newCone']"
-      );
-      await btn.click();
+      await page.getByTitle('New Cone').click();
+
       await page.getByLabel('Radius1').click();
       await page.getByLabel('Radius1').fill('15');
       await page.getByLabel('Radius2').click();