diff --git a/comfy_extras/nodes_v3_test.py b/comfy_extras/nodes_v3_test.py
index 47ab473a9..d6f9fcc76 100644
--- a/comfy_extras/nodes_v3_test.py
+++ b/comfy_extras/nodes_v3_test.py
@@ -14,14 +14,12 @@ class XYZ:
     class Output(io.OutputV3):
         ...
 
-class MyState(io.NodeState):
-    my_str: str
-    my_int: int
-
 
 class V3TestNode(io.ComfyNodeV3):
-
-    state: MyState
+    class State(io.NodeState):
+        my_str: str
+        my_int: int
+    state: State
 
     def __init__(self):
         super().__init__()