From a113c5e3aed805ae2acbca042838ebf897595d3f Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Thu, 12 Jun 2025 17:58:42 +0300
Subject: [PATCH] viz: update browser test to properly shutdown [pr] (#10793)

Using `await page.evaluate` can cause non deterministic `TargetCloseError`
exceptions if it cannot find the elements on the page, Puppeteer
doesn't cleanly stop when `browser.close()` is called.
[Failing CI](https://github.com/tinygrad/tinygrad/actions/runs/15596803685/job/43928961323?pr=10763#step:9:61)
---
 test/web/test_viz.js | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

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();
   }
 }