diff --git a/test/test_axi.py b/test/test_axi.py
index afc6f8ef1..4771aa791 100644
--- a/test/test_axi.py
+++ b/test/test_axi.py
@@ -9,6 +9,8 @@ import random
 
 from migen import *
 
+from litex.gen import *
+
 from litex.soc.interconnect.axi import *
 from litex.soc.interconnect import wishbone
 
@@ -336,49 +338,33 @@ class TestAXI(unittest.TestCase):
             r_ready_random  = 90
         )
 
-    def test_axi_width_converter(self):
-        class DUT(Module):
+    def test_axi_down_converter(self):
+        class DUT(LiteXModule):
             def __init__(self, dw_from=64, dw_to=32):
-                self.axi_master = axi_master = AXIInterface(data_width=dw_from)
-                self.axi_slave  = axi_slave  = AXIInterface(data_width=dw_to)
-                converter = AXIConverter(axi_master, axi_slave)
-                self.submodules += converter
-                wb = wishbone.Interface(data_width=dw_to, address_width=axi_slave.address_width)
-                axi2wb = AXI2Wishbone(axi_slave, wb)
-                self.submodules += axi2wb
-                self.mem = mem = wishbone.SRAM(1024, bus=wb, init=range(256))
-                self.submodules += mem
+                self.axi_master = AXIInterface(data_width=dw_from)
+                axi_slave       = AXIInterface(data_width=dw_to)
+                wb_slave        = wishbone.Interface(data_width=dw_to, address_width=axi_slave.address_width)
+                self.converter = AXIConverter(self.axi_master, axi_slave)
+                self.axi2wb    = AXI2Wishbone(axi_slave, wb_slave)
+                self.mem       = wishbone.SRAM(1024, bus=wb_slave, init=range(256))
 
-        class DUT_ref(Module):
-            """
-            An alternative configuration to the DUT above not using AXIConverter
-            to demonstrate that the generators below are valid.
-            Not used by default.
-            """
-            def __init__(self, dw_from=64, dw_to=32):
-                self.axi_master = axi_master = AXIInterface(data_width=dw_from)
-                wb_from = wishbone.Interface(data_width=dw_from, address_width=axi_master.address_width)
-                axi2wb = AXI2Wishbone(axi_master, wb_from)
-                self.submodules += axi2wb
-                wb_to = wishbone.Interface(data_width=dw_to, address_width=axi_master.address_width)
-                wb2wb = wishbone.Converter(wb_from, wb_to)
-                self.submodules += wb2wb
-                self.mem = mem = wishbone.SRAM(1024, bus=wb_to, init=range(256))
-                self.submodules += mem
-
-        def generator_rd(dut):
+        def read_generator(dut):
             axi_port = dut.axi_master
+
+            # AXI Read.
             addr = 0x34
-            yield axi_port.ar.addr.eq(addr * dut.mem.bus.data_width // 8)
+            yield axi_port.ar.addr.eq(addr * 4)
             yield axi_port.ar.valid.eq(1)
-            yield axi_port.ar.burst.eq(0b1) # CHECKME.
+            yield axi_port.ar.burst.eq(0b1)
             yield axi_port.ar.len.eq(0)
-            yield axi_port.ar.size.eq(log2_int(axi_port.data_width // 8))
+            yield axi_port.ar.size.eq(0b011)
             yield axi_port.r.ready.eq(1)
             yield
             while (yield axi_port.r.valid) == 0:
                 yield
             rd = (yield axi_port.r.data)
+
+            # Check Mem Content.
             mem_content = 0
             i = 0
             while i < axi_port.data_width // dut.mem.bus.data_width:
@@ -386,16 +372,18 @@ class TestAXI(unittest.TestCase):
                 i += 1
             assert rd == mem_content, (hex(rd), hex(mem_content))
 
-        def generator_wr(dut):
+        def write_generator(dut):
             axi_port = dut.axi_master
+
+            # AXI Write.
             addr = 0x24
             data = 0x98761244
             yield axi_port.aw.addr.eq(addr * 4)
             yield axi_port.aw.valid.eq(1)
-            yield axi_port.aw.burst.eq(0b1) # CHECKME.
+            yield axi_port.aw.burst.eq(0b1)
             yield axi_port.aw.len.eq(0)
-            yield axi_port.aw.size.eq(log2_int(axi_port.data_width // 8))
-            yield axi_port.w.strb.eq(2**(len(axi_port.w.data)//8) - 1)
+            yield axi_port.aw.size.eq(0b011)
+            yield axi_port.w.strb.eq(0b111111111)
             yield axi_port.w.data.eq(data)
             yield axi_port.w.valid.eq(1)
             yield axi_port.w.last.eq(1)
@@ -406,6 +394,8 @@ class TestAXI(unittest.TestCase):
             while (yield axi_port.w.ready) == 0:
                 yield
             yield axi_port.w.valid.eq(0)
+
+            # Check Mem Content.
             mem_content = 0
             i = 0
             while i < axi_port.data_width // dut.mem.bus.data_width:
@@ -414,5 +404,4 @@ class TestAXI(unittest.TestCase):
             assert data == mem_content, (hex(data), hex(mem_content))
 
         dut = DUT(64, 32)
-        #dut = DUT_ref(64, 32)
-        run_simulation(dut, [generator_rd(dut), generator_wr(dut)], vcd_name="sim.vcd")
+        run_simulation(dut, [read_generator(dut), write_generator(dut)], vcd_name="sim.vcd")