diff --git a/nodes.py b/nodes.py
index fa957ad..df60b85 100644
--- a/nodes.py
+++ b/nodes.py
@@ -410,7 +410,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "KRestartSamplerCustom": "KSampler With Restarts (Custom)",
 }
 
-if os.environ.get("COMFYUI_RESTART_SAMPLING_SELFTEST", "").strip() == "1":
+if INCLUDE_SELFTEST:
 
     class RestartSelfTestNode:
         @classmethod