diff --git a/tests/axi/test_axi.py b/tests/axi/test_axi.py
index 49ec553..52a52ce 100644
--- a/tests/axi/test_axi.py
+++ b/tests/axi/test_axi.py
@@ -287,7 +287,7 @@ def cycle_pause():
 
 if cocotb.SIM_NAME:
 
-    data_width = int(os.getenv("PARAM_DATA_WIDTH"))
+    data_width = len(cocotb.top.axi_wdata)
     byte_width = data_width // 8
     max_burst_size = (byte_width-1).bit_length()
 
diff --git a/tests/axis/test_axis.py b/tests/axis/test_axis.py
index 6b04e9b..4772342 100644
--- a/tests/axis/test_axis.py
+++ b/tests/axis/test_axis.py
@@ -126,7 +126,7 @@ def cycle_pause():
 
 
 def size_list():
-    data_width = int(os.getenv("PARAM_DATA_WIDTH"))
+    data_width = len(cocotb.top.axis_tdata)
     byte_width = data_width // 8
     return list(range(1, byte_width*4+1)) + [512] + [1]*64