diff --git a/tests-unit/websocket_feature_flags_test.py b/tests-unit/websocket_feature_flags_test.py
index b8d7ee258..f79a3fd19 100644
--- a/tests-unit/websocket_feature_flags_test.py
+++ b/tests-unit/websocket_feature_flags_test.py
@@ -1,4 +1,7 @@
 """Simplified tests for WebSocket feature flags functionality."""
+from comfy.cli_args import args
+# Force CPU mode for unit testing
+args.cpu = True
 from unittest.mock import Mock, patch
 from comfy_api import feature_flags