1 files changed, 18 insertions, 21 deletions
diff --git a/hdl/lib/shift_reg.py b/hdl/lib/shift_reg.py
index 3fb05c1..6966a77 100644
--- a/hdl/lib/shift_reg.py
+++ b/hdl/lib/shift_reg.py
@@ -1,7 +1,8 @@
 from amaranth import *
 from amaranth.sim import Simulator, Settle, Delay
 
-from hdl.utils import cmd
+from hdl.utils import cmd, sim, step
+from hdl.lib.in_out_buff import InOutBuff
 
 class ShiftReg(Elaboratable):
     def __init__(self, width):
@@ -31,52 +32,48 @@ class ShiftReg(Elaboratable):
 
         return m
 
-def test_shift_reg(filename="out.vcd"):
+def test_shiftreg_right():
     dut = ShiftReg(8)
-
-    def proc1():
+    def proc():
         val = 0xAB
 
         yield dut.load_val.eq(val)
         yield dut.en.eq(0)
         yield dut.load.eq(1)
-        yield
-        yield Settle()
+        yield from step()
         yield dut.load.eq(0)
         yield dut.en.eq(1)   
+        yield Settle()
 
         for _ in range(9):
             reg_val = yield dut.reg
             assert reg_val == val, f"Incorrect shift ---EXPECTED: {hex(val)}   ---GOT: {hex(reg_val)}"
             val = val >> 1
-            yield
-            yield Settle()
+            yield from step()
+    sim(dut, proc)
 
+def test_shiftreg_left():
+    dut = ShiftReg(8)
+    def proc():
         val = 0xBD
         yield dut.load_val.eq(val)
+        yield dut.en.eq(0)
         yield dut.load.eq(1)
         yield dut.right_left.eq(1)
-        yield
-        yield Settle()
+        yield from step()
         yield dut.load.eq(0)
+        yield dut.en.eq(1)
+        yield Settle()
 
         for _ in range(9):
             reg_val = yield dut.reg
             assert reg_val == val, f"Incorrect shift ---EXPECTED: {hex(val)}   ---GOT: {hex(reg_val)}"
             val = (val << 1) & 0xff
-            yield
-            yield Settle()
-
-
+            yield from step()
+    sim(dut, proc)
 
-    sim = Simulator(dut)
-    sim.add_clock(1e-6)
-    sim.add_sync_process(proc1)
-    
-    with sim.write_vcd(filename):
-        sim.run()
 
 
 if __name__ == '__main__':
     shift_reg = ShiftReg(8)
-    cmd(shift_reg, test_shift_reg)
-\ No newline at end of file
+    cmd(shift_reg)
+\ No newline at end of file