diff --git a/transactron/testing/gtkw_extension.py b/transactron/testing/gtkw_extension.py
index b13fd0ba2..db407ac9d 100644
--- a/transactron/testing/gtkw_extension.py
+++ b/transactron/testing/gtkw_extension.py
@@ -8,10 +8,8 @@
 
 
 class _VCDWriterExt(_VCDWriter):
-    def __init__(self, fragment, *, vcd_file, gtkw_file, traces):
-        super().__init__(
-            fragment=fragment, vcd_file=vcd_file, gtkw_file=gtkw_file, traces=list(flatten_signals(traces))
-        )
+    def __init__(self, design, *, vcd_file, gtkw_file, traces):
+        super().__init__(design=design, vcd_file=vcd_file, gtkw_file=gtkw_file, traces=list(flatten_signals(traces)))
         self._tree_traces = traces
 
     def close(self, timestamp):
@@ -62,7 +60,7 @@ def gtkw_traces(traces):
 
 @contextmanager
 def write_vcd_ext(engine, vcd_file, gtkw_file, traces):
-    vcd_writer = _VCDWriterExt(engine._fragment, vcd_file=vcd_file, gtkw_file=gtkw_file, traces=traces)
+    vcd_writer = _VCDWriterExt(engine._design, vcd_file=vcd_file, gtkw_file=gtkw_file, traces=traces)
     try:
         engine._vcd_writers.append(vcd_writer)
         yield Tick()
diff --git a/transactron/testing/infrastructure.py b/transactron/testing/infrastructure.py
index 85a09b40c..861428fc3 100644
--- a/transactron/testing/infrastructure.py
+++ b/transactron/testing/infrastructure.py
@@ -184,7 +184,7 @@ def __init__(
             # TODO: try to merge with Amaranth.
             if isinstance(extra_signals, Callable):
                 extra_signals = extra_signals()
-            clocks = [d.clk for d in cast(Any, self)._fragment.domains.values()]
+            clocks = [d.clk for d in cast(Any, self)._design.fragment.domains.values()]
 
             self.ctx = write_vcd_ext(
                 cast(Any, self)._engine,