diff --git a/test/integration/annotation_spec.mjs b/test/integration/annotation_spec.mjs
index 7a6b123e3..579ee2d0c 100644
--- a/test/integration/annotation_spec.mjs
+++ b/test/integration/annotation_spec.mjs
@@ -74,13 +74,7 @@ describe("Annotation highlight", () => {
         pages.map(async ([browserName, page]) => {
           for (const i of [23, 22, 14]) {
             await page.click(`[data-annotation-id='${i}R']`);
-            await page.waitForFunction(
-              id =>
-                document.activeElement ===
-                document.querySelector(`#pdfjs_internal_id_${id}R`),
-              {},
-              i
-            );
+            await page.waitForSelector(`#pdfjs_internal_id_${i}R:focus`);
           }
         })
       );
diff --git a/test/integration/highlight_editor_spec.mjs b/test/integration/highlight_editor_spec.mjs
index 793250926..f506279ea 100644
--- a/test/integration/highlight_editor_spec.mjs
+++ b/test/integration/highlight_editor_spec.mjs
@@ -1591,18 +1591,12 @@ describe("Highlight Editor", () => {
           await page.focus(getEditorSelector(1));
 
           await kbFocusPrevious(page);
-          await page.waitForFunction(
-            sel => document.querySelector(sel) === document.activeElement,
-            {},
-            `.page[data-page-number="1"] > .textLayer`
+          await page.waitForSelector(
+            `.page[data-page-number="1"] > .textLayer:focus`
           );
 
           await kbFocusNext(page);
-          await page.waitForFunction(
-            sel => document.querySelector(sel) === document.activeElement,
-            {},
-            getEditorSelector(1)
-          );
+          await page.waitForSelector(`${getEditorSelector(1)}:focus`);
         })
       );
     });