1 files changed, 20 insertions, 12 deletions
diff --git a/tests/wpt/web-platform-tests/common/worklet-reftest.js b/tests/wpt/web-platform-tests/common/worklet-reftest.js
index 49e129ee0aa..abdda5b05e7 100644
--- a/tests/wpt/web-platform-tests/common/worklet-reftest.js
+++ b/tests/wpt/web-platform-tests/common/worklet-reftest.js
@@ -1,25 +1,33 @@
+// Imports code into a worklet. E.g.
+//
+// importWorklet(CSS.paintWorklet, {url: 'script.js'});
+// importWorklet(CSS.paintWorklet, '/* javascript string */');
+function importWorklet(worklet, code) {
+    let url;
+    if (typeof code === 'object') {
+      url = code.url;
+    } else {
+      const blob = new Blob([code], {type: 'text/javascript'});
+      url = URL.createObjectURL(blob);
+    }
+
+    return worklet.addModule(url);
+}
+
 // To make sure that we take the snapshot at the right time, we do double
 // requestAnimationFrame. In the second frame, we take a screenshot, that makes
 // sure that we already have a full frame.
-function importWorkletAndTerminateTestAfterAsyncPaint(worklet, code) {
+async function importWorkletAndTerminateTestAfterAsyncPaint(worklet, code) {
     if (typeof worklet === 'undefined') {
         takeScreenshot();
         return;
     }
 
-    let url;
-    if (typeof code === 'object') {
-      url = code.url;
-    } else {
-      const blob = new Blob([code], {type: 'text/javascript'});
-      url = URL.createObjectURL(blob);
-    }
+    await importWorklet(worklet, code);
 
-    worklet.addModule(url).then(function() {
+    requestAnimationFrame(function() {
         requestAnimationFrame(function() {
-            requestAnimationFrame(function() {
-                takeScreenshot();
-            });
+            takeScreenshot();
         });
     });
 }