diff --git a/test/common/functions.py b/test/common/functions.py
index 6c7fa4412..6a11e41ab 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -1,10 +1,11 @@
 from amaranth import *
-from typing import TYPE_CHECKING
+from amaranth.hdl.ast import Statement
+from amaranth.sim.core import Command
+from typing import TypeVar, Any, Generator, TypeAlias
 from coreblocks.utils._typing import RecordValueDict, RecordIntDict
 
-
-if TYPE_CHECKING:
-    from .testbenchio import TestGen
+T = TypeVar("T")
+TestGen: TypeAlias = Generator[Command | Value | Statement | None, Any, T]
 
 
 def set_inputs(values: RecordValueDict, field: Record) -> "TestGen[None]":
diff --git a/test/common/testbenchio.py b/test/common/testbenchio.py
index c1101cb79..b8ccd1ecb 100644
--- a/test/common/testbenchio.py
+++ b/test/common/testbenchio.py
@@ -1,16 +1,11 @@
 from amaranth import *
-from amaranth.hdl.ast import Statement
-from amaranth.sim.core import Command
 from amaranth.sim import Settle, Passive
-from typing import Any, Generator, TypeVar, Optional, Callable
+from typing import Optional, Callable
 from transactron.lib import AdapterBase
 from transactron.core import ValueLike, SignalBundle
 from transactron._utils import method_def_helper
 from coreblocks.utils._typing import RecordIntDictRet, RecordValueDict, RecordIntDict
-from .functions import set_inputs, get_outputs
-
-T = TypeVar("T")
-TestGen = Generator[Command | Value | Statement | None, Any, T]
+from .functions import set_inputs, get_outputs, TestGen
 
 
 class TestbenchIO(Elaboratable):