diff --git a/tests/silimate/opt_balance_tree.ys b/tests/silimate/opt_balance_tree.ys
index ed683cb9b..8427d25cb 100644
--- a/tests/silimate/opt_balance_tree.ys
+++ b/tests/silimate/opt_balance_tree.ys
@@ -1,3 +1,44 @@
+log -header "Should be turned into a tree"
+log -push
+design -reset
+read_verilog <<EOF
+module top (
+  input wire a,
+  input wire b,
+  input wire c,
+  input wire d,
+  output wire x,
+);
+  assign x = a & b & c & d;
+endmodule
+EOF
+check -assert
+
+autoname
+write_json pre.json
+exec -- netlistsvg pre.json -o pre.svg
+
+# Check equivalence after opt_balance_tree
+equiv_opt -assert opt_balance_tree
+
+design -load postopt
+# opt_balance_tree
+
+autoname
+write_json post.json
+exec -- netlistsvg post.json -o post.svg
+
+# Checks if y is still wired up to the correct gate
+# select -set y_wires o:y %ci
+# select -set y_driver @y_wires %ci
+# select -set and_y_cell t:$and @y_driver %i
+# select @and_y_cell -assert-count 1
+# select -set inputs @and_y_cell %ci
+# select -assert-count 1 @inputs i:c %i
+
+design -reset
+log -pop
+
 log -header "Should not be turned into a tree"
 log -push
 design -reset
@@ -18,11 +59,18 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# design -load postopt
+design -load postopt
+
+# Checks if y is still wired up to the correct gate
+select -set y_wires o:y %ci
+select -set y_driver @y_wires %ci
+select -set and_y_cell t:$and @y_driver %i
+select @and_y_cell -assert-count 1
+select -set inputs @and_y_cell %ci
+select -assert-count 1 @inputs i:c %i
 
 design -reset
 log -pop
@@ -49,12 +97,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
@@ -77,12 +124,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
@@ -108,12 +154,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
@@ -137,12 +182,11 @@ endmodule
 EOF
 check -assert
 
-# Check equivalence after opt_expand
+# Check equivalence after opt_balance_tree
 equiv_opt -assert opt_balance_tree
 
-# TODO check design to make sure that it is not turned into a tree
-# Note: this one already worked, just including here for completeness
-# design -load postopt
+design -load postopt
+select -assert-count 3 t:$and
 
 design -reset
 log -pop
\ No newline at end of file