From afc3904e589161f19205c4bde4c76f6fbc62ab6d Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Fri, 17 Apr 2026 16:34:44 +0300
Subject: [PATCH] viz/cli: unit tests in CI (#15788)

* simple failing test

* test stdout

* cleanup sqttmap
---
 extra/viz/cli.py         |  4 ++--
 test/amd/test_sqttmap.py | 10 ++--------
 test/null/test_viz.py    | 30 +++++++++++++++++++++++++++++-
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/extra/viz/cli.py b/extra/viz/cli.py
index 94718418cc..eb56089b9e 100755
--- a/extra/viz/cli.py
+++ b/extra/viz/cli.py
@@ -165,8 +165,8 @@ def main(args) -> None:
         print(f"{fmt_colored(name)}{' ' * max(0, 36 - ansilen(name))} {time_to_str(t, w=9)} {c:7d} {t/total*100.0:6.2f}%")
         if ref is not None:
           steps = rewrites[viz_data.ctxs[ref]["name"]]
-          if DEBUG >= 3: print_step(get(steps, "View Base AST"))
-          if DEBUG >= 4: print_step(get(steps, "View Source"))
+          if DEBUG >= 3 and (ast_step:=steps.get("View Base AST")) is not None: print_step(ast_step)
+          if DEBUG >= 4: print_step(steps["View Source"])
       if num_rows > 0 and items[num_rows:]:
         other_t = sum(t for _,(t,_,_) in items[num_rows:])
         other_c = sum(c for _,(_,c,_) in items[num_rows:])
diff --git a/test/amd/test_sqttmap.py b/test/amd/test_sqttmap.py
index d6225e61fa..869befa2b7 100644
--- a/test/amd/test_sqttmap.py
+++ b/test/amd/test_sqttmap.py
@@ -1,5 +1,5 @@
 # test to compare every packet with the rocprof decoder
-import unittest, pickle, contextlib, io
+import unittest, pickle
 from typing import Iterator
 from pathlib import Path
 from tinygrad.helpers import DEBUG, getenv, temp, ansistrip
@@ -7,17 +7,11 @@ from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
 from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm
+from test.null.test_viz import run_cli
 
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
 
-def run_cli(*cli_args) -> str:
-  from extra.viz.cli import main, get_arg_parser
-  args = get_arg_parser().parse_args(cli_args)
-  with contextlib.redirect_stdout(buf:=io.StringIO()):
-    main(args)
-  return buf.getvalue().strip()
-
 def rocprof_inst_traces_match(sqtt, prg, target):
   from tinygrad.viz.serve import amd_decode
   from extra.sqtt.roc import decode as roc_decode, InstExec
diff --git a/test/null/test_viz.py b/test/null/test_viz.py
index e98f6b4458..0d8f15cf23 100644
--- a/test/null/test_viz.py
+++ b/test/null/test_viz.py
@@ -1,4 +1,5 @@
-import unittest, decimal, sys, json, contextlib
+import unittest, decimal, sys, json, contextlib, tempfile, pickle, io
+from pathlib import Path
 from dataclasses import dataclass
 from typing import Generator
 
@@ -884,5 +885,32 @@ class TestCfg(unittest.TestCase):
     k.emit(s_code_end())
     self.get_cfg("jump_back_to_end", k)
 
+# launch viz cli without subprocess
+def run_cli(*cli_args) -> str:
+  from extra.viz.cli import main, get_arg_parser
+  args = get_arg_parser().parse_args(cli_args)
+  with contextlib.redirect_stdout(buf:=io.StringIO()):
+    main(args)
+  return buf.getvalue().strip()
+
+class TestCLI(unittest.TestCase):
+  def test_simple(self):
+    a = Tensor.empty(1, device="NULL")+2.0
+    def custom_empty_prg(B:UOp, A:UOp) -> UOp:
+      sink = UOp(Ops.SINK, arg=KernelInfo(name="custom_empty"))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=a.device), UOp(Ops.LINEAR, src=(sink,))))
+    b = Tensor.custom_kernel(Tensor.empty_like(a), a, fxn=custom_empty_prg)[0]
+    with save_viz() as viz:
+      b.realize()
+    # save trace to disk for CLI to consume it
+    with tempfile.TemporaryDirectory() as tmpdir:
+      (r:=Path(tmpdir)/"rewrites.pkl").write_bytes(pickle.dumps(viz.data.trace))
+      (p:=Path(tmpdir)/"profile.pkl").write_bytes(pickle.dumps(cpu_events))
+      with Context(DEBUG=4):
+        kernels = run_cli("--rewrites-path", str(r), "--profile-path", str(p), "-p", "-s", "NULL")
+      self.assertIn("void custom_empty", kernels)
+      self.assertIn("E", kernels)
+      self.assertIn("UOp.const", kernels)
+
 if __name__ == "__main__":
   unittest.main()