diff --git a/tests/browser/test/procedure_test.js b/tests/browser/test/procedure_test.js
index 2d7489a18..69e66c7c1 100644
--- a/tests/browser/test/procedure_test.js
+++ b/tests/browser/test/procedure_test.js
@@ -31,6 +31,7 @@ suite('Testing Connecting Blocks', function (done) {
   test('Testing Procedure', async function () {
     // Drag out first function
     let proceduresDefReturn = await getBlockTypeFromCategory(
+      browser,
       'Functions',
       'procedures_defreturn'
     );
@@ -39,6 +40,7 @@ suite('Testing Connecting Blocks', function (done) {
 
     // Drag out second function.
     proceduresDefReturn = await getBlockTypeFromCategory(
+      browser,
       'Functions',
       'procedures_defreturn'
     );
@@ -46,7 +48,11 @@ suite('Testing Connecting Blocks', function (done) {
     await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
 
     // Drag out numeric
-    const mathNumeric = await getBlockTypeFromCategory('Math', 'math_number');
+    const mathNumeric = await getBlockTypeFromCategory(
+      browser,
+      'Math',
+      'math_number'
+    );
     await mathNumeric.dragAndDrop({x: 50, y: 20});
     await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
 
@@ -64,17 +70,29 @@ suite('Testing Connecting Blocks', function (done) {
     const doSomething2 = await browser.$(
       '#content_blocks > div > svg.blocklySvg > g > g.blocklyBlockCanvas > g:nth-child(2)'
     );
-    const doSomethingFlyout = await getNthBlockOfCategory('Functions', 3);
+    const doSomethingFlyout = await getNthBlockOfCategory(
+      browser,
+      'Functions',
+      3
+    );
     await doSomethingFlyout.dragAndDrop(doSomething2);
     await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
     const doSomethingFlyoutWorkspace = await getSelectedBlockElement(browser);
     await doSomethingFlyoutWorkspace.dragAndDrop({x: 130, y: 20});
 
     // Drag out print from flyout and connect it with doSomething 2
-    const printFlyout = await getBlockTypeFromCategory('Text', 'text_print');
+    const printFlyout = await getBlockTypeFromCategory(
+      browser,
+      'Text',
+      'text_print'
+    );
     await printFlyout.dragAndDrop({x: 50, y: 20});
     await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
-    const doSomething2Flyout = await getNthBlockOfCategory('Functions', 4);
+    const doSomething2Flyout = await getNthBlockOfCategory(
+      browser,
+      'Functions',
+      4
+    );
     await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
     await doSomething2Flyout.dragAndDrop({x: 130, y: 20});
     await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
diff --git a/tests/browser/test/test_setup.js b/tests/browser/test/test_setup.js
index 3b15bb5f0..819ae7d92 100644
--- a/tests/browser/test/test_setup.js
+++ b/tests/browser/test/test_setup.js
@@ -95,7 +95,7 @@ async function getSelectedBlockElement(browser) {
 async function getBlockElementById(browser, id) {
   return await browser.$(`[data-id="${id}"]`);
 }
-async function getCategory(categoryName) {
+async function getCategory(browser, categoryName) {
   const categories = await browser.$$('.blocklyTreeLabel');
 
   let category;
@@ -110,8 +110,8 @@ async function getCategory(categoryName) {
   return category;
 }
 
-async function getNthBlockOfCategory(categoryName, n) {
-  const category = await getCategory(categoryName);
+async function getNthBlockOfCategory(browser, categoryName, n) {
+  const category = await getCategory(browser, categoryName);
   category.click();
   await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
   const block = await browser.$(
@@ -120,17 +120,17 @@ async function getNthBlockOfCategory(categoryName, n) {
   return block;
 }
 
-async function getBlockTypeFromCategory(categoryName, blockType) {
-  const category = await getCategory(categoryName);
+async function getBlockTypeFromCategory(browser, categoryName, blockType) {
+  const category = await getCategory(browser, categoryName);
   category.click();
   await new Promise((resolve) => setTimeout(resolve, 2000)); // 2 sec
 
-  const id = browser.execute(() => {
+  const id = await browser.execute((blockType) => {
     return Blockly.getMainWorkspace()
       .getFlyout()
       .getWorkspace()
       .getBlocksByType(blockType)[0].id;
-  });
+  }, blockType);
   return await browser.$(`[data-id="${id}"]`);
 }