diff --git a/cz.py b/cz.py
index 5f607bf..8591dd3 100644
--- a/cz.py
+++ b/cz.py
@@ -1,6 +1,5 @@
-from matplotlib import pyplot as plt
 from graph import *
-from viz import draw
+import viz
 import itertools as it
 import clifford
 
@@ -16,15 +15,12 @@ def local_complementation(g, vops, v):
         vops[i] = clifford.times_table[vops[i]][clifford.msqz]
 
 if __name__ == '__main__':
-    g, vops = graph(10)
+    g, vops = graph()
     add_edge(g, 0, 1)
-    add_edge(g, 1, 3)
-    add_edge(g, 3, 2)
-    add_edge(g, 3, 0)
-    add_edge(g, 2, 0)
-    add_edge(g, 0, 5)
-    vops[0] = 1
-    draw(g, vops)
-    plt.clf()
+    add_edge(g, 1, 2)
+    add_edge(g, 0, 2)
+    add_edge(g, 0, 3)
+
+    viz.draw(g, vops, "out.pdf")
     local_complementation(g, vops, 0)
-    draw(g, vops, "out2.pdf")
+    viz.draw(g, vops, "out2.pdf")
diff --git a/tests/test_cz.py b/tests/test_cz.py
new file mode 100644
index 0000000..777c3c4
--- /dev/null
+++ b/tests/test_cz.py
@@ -0,0 +1,19 @@
+import cz
+from graph import *
+import viz
+
+def test_local_complementation():
+    """ Test that local complementation works as expected """
+    g, vops = graph()
+    add_edge(g, 0, 1)
+    add_edge(g, 0, 2)
+    add_edge(g, 1, 2)
+    add_edge(g, 0, 3)
+    cz.local_complementation(g, vops, 0)
+    assert has_edge(g, 0, 1)
+    assert has_edge(g, 0, 2)
+    assert not has_edge(g, 1, 2)
+    assert has_edge(g, 3, 2)
+    assert has_edge(g, 3, 1)
+
+    # TODO: test VOP conditions
diff --git a/viz.py b/viz.py
index 3bbafeb..edd9d24 100644
--- a/viz.py
+++ b/viz.py
@@ -6,10 +6,11 @@ import networkx as nx
 from matplotlib import pyplot as plt
 from graph import *
 
-VOP_COLORS = ["red", "green", "blue"]
+VOP_COLORS = ["red", "green", "blue", "orange", "yellow", "purple", "black", "white"]
 
 def draw(graph, vops, filename="out.pdf", ns=500):
     """ Draw a graph with networkx layout """
+    plt.clf()
     g = nx.from_edgelist(edgelist(graph))
     pos = nx.spring_layout(g)
     colors = [VOP_COLORS[vop % len(VOP_COLORS)] for vop in vops]
@@ -21,7 +22,7 @@ def draw(graph, vops, filename="out.pdf", ns=500):
     plt.savefig(filename)
 
 if __name__ == '__main__':
-    g, vops = graph(10)
+    g, vops = graph()
     add_edge(g, 0, 1)
     add_edge(g, 1, 3)
     add_edge(g, 3, 2)