diff --git a/test/web/test_viz.js b/test/web/test_viz.js
index 80475734b9..d39bc0781b 100644
--- a/test/web/test_viz.js
+++ b/test/web/test_viz.js
@@ -10,23 +10,24 @@ async function main() {
   }));
 
   // ** run browser tests
-  let browser;
+  let browser, page;
   try {
     browser = await puppeteer.launch({ headless: true });
-    const page = await browser.newPage();
+    page = await browser.newPage();
     const res = await page.goto("http://localhost:8000", { waitUntil:"domcontentloaded" });
     if (res.status() !== 200) throw new Error("Failed to load page");
     const scheduleSelector = await page.waitForSelector("ul");
     scheduleSelector.click();
     await page.waitForSelector("rect");
-    const nodes = await page.evaluate(() => document.querySelectorAll("#nodes > g").length);
-    const edges = await page.evaluate(() => document.querySelectorAll("#edges > path").length);
-    if (!nodes || !edges) {
-      throw new Error("VIZ didn't render a graph")
-    }
+    await page.waitForFunction(() => {
+      const nodes = document.querySelectorAll("#nodes > g").length;
+      const edges = document.querySelectorAll("#edges > path").length;
+      return nodes > 0 && edges > 0;
+    });
   } finally {
     // ** cleanups
-    if (browser) await browser.close();
+    if (page != null) await page.close();
+    if (browser != null) await browser.close();
     proc.kill();
   }
 }