From 7114b6ab3198c9536f7791aedd6f736c39dadcff Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Wed, 4 Jun 2025 14:58:24 +0300
Subject: [PATCH] viz browser tests (#10626)

* viz browser tests

* expect failure if js/ isn't included

* back green
---
 .github/workflows/test.yml | 14 +++++++++-----
 test/web/test_viz.js       | 34 ++++++++++++++++++++++++++++++++++
 tinygrad/viz/serve.py      |  2 +-
 3 files changed, 44 insertions(+), 6 deletions(-)
 create mode 100644 test/web/test_viz.js

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 9a3c9ccbef..cefd235ee3 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -53,11 +53,6 @@ jobs:
         source venv/bin/activate
         pip install $GITHUB_WORKSPACE
         python -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
-        # Test using VIZ=1 after package installation
-        VIZ=1 python -c "from tinygrad.tensor import Tensor; Tensor([1,2,3,4,5]).realize()" & VIZ_PID=$!
-        echo "started VIZ server at: $(ps -p "$VIZ_PID" -o pid=)"
-        i=0; while ((i++ < 10)); do curl -sSf localhost:8000 > /dev/null && break || { echo "VIZ verification attempt $i/10"; sleep 1; }; done; ((i > 10)) && echo "Could not verify VIZ server" && exit 1
-        kill $VIZ_PID
         pip install mypy
         mypy -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
     - name: Run beautiful_mnist with tinygrad only
@@ -784,6 +779,15 @@ jobs:
       run: npm install puppeteer
     - name: Run WEBGPU Efficientnet
       run: node test/web/test_webgpu.js
+    - name: Run VIZ tests as external package
+      run: |
+        mkdir $GITHUB_WORKSPACE/test_dir
+        cd $GITHUB_WORKSPACE/test_dir
+        python -m venv venv
+        source venv/bin/activate
+        pip install $GITHUB_WORKSPACE
+        cp $GITHUB_WORKSPACE/test/web/test_viz.js .
+        node test_viz.js
 
   osxremote:
     name: MacOS (remote metal)
diff --git a/test/web/test_viz.js b/test/web/test_viz.js
new file mode 100644
index 0000000000..69d43bd96c
--- /dev/null
+++ b/test/web/test_viz.js
@@ -0,0 +1,34 @@
+const { spawn } = require("child_process");
+const puppeteer = require("puppeteer");
+
+async function main() {
+  // ** start viz server
+  const proc = spawn("python", ["-u", "-c", "from tinygrad import Tensor; Tensor.arange(4).realize()"], { env: { ...process.env, VIZ:"1" },
+                      stdio: ["inherit", "pipe", "inherit"]});
+  await new Promise(resolve => proc.stdout.on("data", r => {
+    if (r.includes("ready")) resolve();
+  }));
+
+  // ** run browser tests
+  let browser;
+  try {
+    browser = await puppeteer.launch({ headless: true });
+    const page = await browser.newPage();
+    const res = await page.goto("http://localhost:8000");
+    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")
+    }
+  } finally {
+    // ** cleanups
+    if (browser) await browser.close();
+    proc.kill();
+  }
+}
+
+main();
diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py
index 273cf4044a..f15aa45252 100755
--- a/tinygrad/viz/serve.py
+++ b/tinygrad/viz/serve.py
@@ -208,7 +208,7 @@ if __name__ == "__main__":
   reloader_thread = threading.Thread(target=reloader)
   reloader_thread.start()
   print(f"*** started viz on {HOST}:{PORT}")
-  print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"))
+  print(colored(f"*** ready in {(time.perf_counter()-st)*1e3:4.2f}ms", "green"), flush=True)
   if len(getenv("BROWSER", "")) > 0: webbrowser.open(f"{HOST}:{PORT}{'/profiler' if contexts is None else ''}")
   try: server.serve_forever()
   except KeyboardInterrupt: