diff --git a/tests/lu_tb.sv b/tests/lu_tb.sv
index a5a21bb..978f70a 100644
--- a/tests/lu_tb.sv
+++ b/tests/lu_tb.sv
@@ -7,30 +7,30 @@ module lu_tb #(parameter W=4)();
     logic [W-1:0] C;
     logic [W-1:0] C_exp[16];
 
-    lu #(W/2) dut();
+    lu #(W/2) dut(op, A, B, C);
 
     integer error_count;
     initial begin
         // AND
-        C_exp[0]  = {W{1'b0}};
-        C_exp[1]  = {W{1'b0}};
-        C_exp[2]  = {W{1'b0}};
-        C_exp[3]  = {W{1'b0}};
+        C_exp[0]  = {{W{1'b0}}, {1'b0}};
+        C_exp[1]  = {{W{1'b0}}, {1'b0}};
+        C_exp[2]  = {{W{1'b0}}, {1'b0}};
+        C_exp[3]  = {{W{1'b0}}, {1'b1}};
         // OR
-        C_exp[4]  = {W{1'b0}};
-        C_exp[5]  = {W{1'b1}};
-        C_exp[6]  = {W{1'b1}};
-        C_exp[7]  = {W{1'b1}};
+        C_exp[4]  = {{W{1'b0}}, {1'b0}};
+        C_exp[5]  = {{W{1'b0}}, {1'b1}};
+        C_exp[6]  = {{W{1'b0}}, {1'b1}};
+        C_exp[7]  = {{W{1'b0}}, {1'b1}};
         // XOR
-        C_exp[8]  = {W{1'b0}};
-        C_exp[9]  = {W{1'b1}};
-        C_exp[10] = {W{1'b1}};
-        C_exp[11] = {W{1'b0}};
+        C_exp[8]  = {{W{1'b0}}, {1'b0}};
+        C_exp[9]  = {{W{1'b0}}, {1'b1}};
+        C_exp[10] = {{W{1'b0}}, {1'b1}};
+        C_exp[11] = {{W{1'b0}}, {1'b0}};
         // NOR
-        C_exp[12] = {W{1'b1}};
-        C_exp[13] = {W{1'b0}};
-        C_exp[14] = {W{1'b0}};
-        C_exp[15] = {W{1'b0}};
+        C_exp[12] = {{W{1'b1}}, {1'b1}};
+        C_exp[13] = {{W{1'b1}}, {1'b0}};
+        C_exp[14] = {{W{1'b1}}, {1'b0}};
+        C_exp[15] = {{W{1'b1}}, {1'b0}};
 
         error_count = 0;
         op = 0;
@@ -42,12 +42,13 @@ module lu_tb #(parameter W=4)();
                 for (integer K = 0; K < 2; ++K) begin
                     #1;
                     if (C != C_exp[N*2*2 + M*2 + K]) begin
-                        $display("\t\x1b[38;5;9mOp 0b%b(0b%b, 0b%b) produced 0b%b but 0b%b was expected\x1b[m", op, A, B, C, C_exp[N*2*2 + M*2 + K]);
+                        $display("\t\x1b[38;5;9mLU(%0d) Op 0b%b(0b%b, 0b%b) produced 0b%b but 0b%b was expected\x1b[m", N*2*2+M*2+K, op, A, B, C, C_exp[N*2*2 + M*2 + K]);
                         error_count = error_count + 1;
                     end
                     B = 1;
                 end
                 A = 1;
+                B = 0;
             end
             A = 0;
             B = 0;