Throw error if pygraphviz not present

Also add assert to check for image and xfail
test if graphviz is not available.

shark_turbine/kernel/wave/visualization.py

@@ -1,16 +1,15 @@
-disabled_graphviz = False
+graphviz_disabled = False
 try:
     import pygraphviz as pgv
 except:
-    disabled_graphviz = True
+    graphviz_disabled = True
 from torch import fx
 import warnings
 
 
 def visualize_graph(graph: fx.Graph, file_name: str):
-    if disabled_graphviz:
-        warnings.warn("pygraphviz not installed, skipping visualization.")
-        return
+    if graphviz_disabled:
+        raise ImportError("pygraphviz not installed, cannot visualize graph")
     G = pgv.AGraph(directed=True)
     for node in graph.nodes:
         G.add_node(node.name)

tests/kernel/wave/visualization_test.py

@@ -1,6 +1,8 @@
 import logging
 from typing import Callable
 import unittest
+import os
+import pytest
 import shark_turbine.kernel as tk
 import shark_turbine.kernel.lang as tkl
 import shark_turbine.kernel.wave as tkw
@@ -54,6 +56,14 @@ def repeat(acc: tkl.Register[M, N, tkl.f32]) -> tkl.Register[M, N, tkl.f32]:
     tkw.write(repeat, c, elements_per_thread=4)
 
 
+graphviz_disabled = False
+try:
+    import pygraphviz
+except:
+    graphviz_disabled = True
+
+
+@pytest.mark.xfail(condition=graphviz_disabled, reason="pygraphviz not installed")
 @run
 def test_gemm():
     constraints: list[tkw.Constraint] = [tkw.WorkgroupConstraint(M, BLOCK_M, 0)]
@@ -75,6 +85,7 @@ def test_gemm():
         IndexingContext.current().finalize()
         expand_graph(graph, constraints)
         visualize_graph(graph.get_subgraph("region_0"), "gemm.png")
+        assert os.path.exists("gemm.png")
 
 
 if __name__ == "__main__":