mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-07-27 16:26:39 +00:00
2b653e8c18
* Support for async execution functions This commit adds support for node execution functions defined as async. When a node's execution function is defined as async, we can continue executing other nodes while it is processing. Standard uses of `await` should "just work", but people will still have to be careful if they spawn actual threads. Because torch doesn't really have async/await versions of functions, this won't particularly help with most locally-executing nodes, but it does work for e.g. web requests to other machines. In addition to the execute function, the `VALIDATE_INPUTS` and `check_lazy_status` functions can also be defined as async, though we'll only resolve one node at a time right now for those. * Add the execution model tests to CI * Add a missing file It looks like this got caught by .gitignore? There's probably a better place to put it, but I'm not sure what that is. * Add the websocket library for automated tests * Add additional tests for async error cases Also fixes one bug that was found when an async function throws an error after being scheduled on a task. * Add a feature flags message to reduce bandwidth We now only send 1 preview message of the latest type the client can support. We'll add a console warning when the client fails to send a feature flags message at some point in the future. * Add async tests to CI * Don't actually add new tests in this PR Will do it in a separate PR * Resolve unit test in GPU-less runner * Just remove the tests that GHA can't handle * Change line endings to UNIX-style * Avoid loading model_management.py so early Because model_management.py has a top-level `logging.info`, we have to be careful not to import that file before we call `setup_logging`. If we do, we end up having the default logging handler registered in addition to our custom one.
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
"""Simplified tests for WebSocket feature flags functionality."""
|
|
from comfy_api import feature_flags
|
|
|
|
|
|
class TestWebSocketFeatureFlags:
|
|
"""Test suite for WebSocket feature flags integration."""
|
|
|
|
def test_server_feature_flags_response(self):
|
|
"""Test server feature flags are properly formatted."""
|
|
features = feature_flags.get_server_features()
|
|
|
|
# Check expected server features
|
|
assert "supports_preview_metadata" in features
|
|
assert features["supports_preview_metadata"] is True
|
|
assert "max_upload_size" in features
|
|
assert isinstance(features["max_upload_size"], (int, float))
|
|
|
|
def test_progress_py_checks_feature_flags(self):
|
|
"""Test that progress.py checks feature flags before sending metadata."""
|
|
# This simulates the check in progress.py
|
|
client_id = "test_client"
|
|
sockets_metadata = {"test_client": {"feature_flags": {}}}
|
|
|
|
# The actual check would be in progress.py
|
|
supports_metadata = feature_flags.supports_feature(
|
|
sockets_metadata, client_id, "supports_preview_metadata"
|
|
)
|
|
|
|
assert supports_metadata is False
|
|
|
|
def test_multiple_clients_different_features(self):
|
|
"""Test handling multiple clients with different feature support."""
|
|
sockets_metadata = {
|
|
"modern_client": {
|
|
"feature_flags": {"supports_preview_metadata": True}
|
|
},
|
|
"legacy_client": {
|
|
"feature_flags": {}
|
|
}
|
|
}
|
|
|
|
# Check modern client
|
|
assert feature_flags.supports_feature(
|
|
sockets_metadata, "modern_client", "supports_preview_metadata"
|
|
) is True
|
|
|
|
# Check legacy client
|
|
assert feature_flags.supports_feature(
|
|
sockets_metadata, "legacy_client", "supports_preview_metadata"
|
|
) is False
|
|
|
|
def test_feature_negotiation_message_format(self):
|
|
"""Test the format of feature negotiation messages."""
|
|
# Client message format
|
|
client_message = {
|
|
"type": "feature_flags",
|
|
"data": {
|
|
"supports_preview_metadata": True,
|
|
"api_version": "1.0.0"
|
|
}
|
|
}
|
|
|
|
# Verify structure
|
|
assert client_message["type"] == "feature_flags"
|
|
assert "supports_preview_metadata" in client_message["data"]
|
|
|
|
# Server response format (what would be sent)
|
|
server_features = feature_flags.get_server_features()
|
|
server_message = {
|
|
"type": "feature_flags",
|
|
"data": server_features
|
|
}
|
|
|
|
# Verify structure
|
|
assert server_message["type"] == "feature_flags"
|
|
assert "supports_preview_metadata" in server_message["data"]
|
|
assert server_message["data"]["supports_preview_metadata"] is True
|