mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-03 00:57:05 +08:00
bd78a3455f
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.
129 lines
4.5 KiB
Python
129 lines
4.5 KiB
Python
"""Simplified tests for WebSocket feature flags functionality."""
|
|
from unittest.mock import Mock, patch
|
|
from comfy_api import feature_flags
|
|
|
|
|
|
class TestWebSocketFeatureFlags:
|
|
"""Test suite for WebSocket feature flags integration."""
|
|
|
|
@patch('main.server')
|
|
def test_preview_message_with_feature_support(self, mock_server):
|
|
"""Test that UNENCODED_PREVIEW_IMAGE is not sent when client supports metadata."""
|
|
# Setup mock server with client that supports preview metadata
|
|
mock_server.client_id = "test_client"
|
|
mock_server.sockets_metadata = {
|
|
"test_client": {
|
|
"feature_flags": {"supports_preview_metadata": True}
|
|
}
|
|
}
|
|
mock_server.send_sync = Mock()
|
|
|
|
# Import the function we're testing
|
|
from main import hijack_progress
|
|
|
|
# Call hijack_progress
|
|
hijack_progress(mock_server)
|
|
|
|
# Verify feature flags can check support
|
|
assert feature_flags.supports_feature(
|
|
mock_server.sockets_metadata,
|
|
"test_client",
|
|
"supports_preview_metadata"
|
|
) is True
|
|
|
|
@patch('main.server')
|
|
def test_preview_message_without_feature_support(self, mock_server):
|
|
"""Test that UNENCODED_PREVIEW_IMAGE is sent when client doesn't support metadata."""
|
|
# Setup mock server with legacy client
|
|
mock_server.client_id = "legacy_client"
|
|
mock_server.sockets_metadata = {
|
|
"legacy_client": {
|
|
"feature_flags": {} # No features
|
|
}
|
|
}
|
|
mock_server.send_sync = Mock()
|
|
|
|
# Import the function we're testing
|
|
from main import hijack_progress
|
|
|
|
# Call hijack_progress
|
|
hijack_progress(mock_server)
|
|
|
|
# Verify feature flags check returns False
|
|
assert feature_flags.supports_feature(
|
|
mock_server.sockets_metadata,
|
|
"legacy_client",
|
|
"supports_preview_metadata"
|
|
) is False
|
|
|
|
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
|