From 77543d3c188793613088b99b6d5aed31827657b1 Mon Sep 17 00:00:00 2001
From: Erik Pasternak <epastern@google.com>
Date: Fri, 27 Jun 2025 11:44:09 -0700
Subject: [PATCH] Fix tests for opening categories

---
 tests/browser/test/test_setup.mjs        |  1 +
 tests/browser/test/toolbox_drag_test.mjs | 37 +++++++++++++++++++++---
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/tests/browser/test/test_setup.mjs b/tests/browser/test/test_setup.mjs
index edfad4d1b..040ca3df0 100644
--- a/tests/browser/test/test_setup.mjs
+++ b/tests/browser/test/test_setup.mjs
@@ -62,6 +62,7 @@ export async function driverSetup() {
   // Use Selenium to bring up the page
   console.log('Starting webdriverio...');
   driver = await webdriverio.remote(options);
+  driver.setWindowSize(800, 600);
   return driver;
 }
 
diff --git a/tests/browser/test/toolbox_drag_test.mjs b/tests/browser/test/toolbox_drag_test.mjs
index 742872d93..1845336eb 100644
--- a/tests/browser/test/toolbox_drag_test.mjs
+++ b/tests/browser/test/toolbox_drag_test.mjs
@@ -11,6 +11,7 @@
 import * as chai from 'chai';
 import {
   getCategory,
+  getDraggableBlockElementByType,
   PAUSE_TIME,
   screenDirection,
   scrollFlyout,
@@ -81,6 +82,32 @@ async function elementInBounds(browser, element) {
   }, element);
 }
 
+/**
+ * Get the type of the nth block in the specified category.
+ * @param browser The active WebdriverIO Browser object.
+ * @param categoryName The name of the category to inspect.
+ * @param n The index of the block to get
+ * @returns A Promise resolving to the type the block in the specified
+ *     category's flyout at index i.
+ */
+async function getNthBlockType(browser, categoryName, n) {
+  const category = await getCategory(browser, categoryName);
+  await category.click();
+  await browser.pause(PAUSE_TIME);
+
+  const blockType = await browser.execute((i) => {
+    return Blockly.getMainWorkspace()
+      .getFlyout()
+      .getWorkspace()
+      .getTopBlocks(false)[i].type;
+  }, n);
+
+  // Unicode escape to close flyout.
+  await browser.keys(['\uE00C']);
+  await browser.pause(PAUSE_TIME);
+  return blockType;
+}
+
 /**
  * Get how many top-level blocks there are in the specified category.
  * @param browser The active WebdriverIO Browser object.
@@ -145,14 +172,16 @@ async function openCategories(browser, categoryList, directionMultiplier) {
           await browser.pause(PAUSE_TIME);
           continue;
         }
-        const flyoutBlock = await browser.$(
-          `.blocklyFlyout .blocklyBlockCanvas > g:nth-child(${3 + i * 2})`,
+        const blockType = await getNthBlockType(browser, categoryName, i);
+        const flyoutBlock = await getDraggableBlockElementByType(
+          browser,
+          categoryName,
+          blockType,
         );
         while (!(await elementInBounds(browser, flyoutBlock))) {
           await scrollFlyout(browser, 0, 50);
         }
-
-        await flyoutBlock.dragAndDrop({x: directionMultiplier * 50, y: 0});
+        await flyoutBlock.click();
         await browser.pause(PAUSE_TIME);
         // Should be one top level block on the workspace.
         const topBlockCount = await browser.execute(() => {