File size: 11,227 Bytes
04eb75a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
"""Tests for feature flags functionality."""

import pytest

from comfy_api.feature_flags import (
    get_connection_feature,
    supports_feature,
    get_server_features,
    CLI_FEATURE_FLAG_REGISTRY,
    SERVER_FEATURE_FLAGS,
    _coerce_flag_value,
    _parse_cli_feature_flags,
)
from comfy.comfy_api_env import (
    environment_overrides_for_base,
    get_environment_overrides,
    normalize_comfy_api_base,
)


class TestFeatureFlags:
    """Test suite for feature flags functions."""

    def test_get_server_features_returns_copy(self):
        """Test that get_server_features returns a copy of the server flags."""
        features = get_server_features()
        # Verify it's a copy by modifying it
        features["test_flag"] = True
        # Original should be unchanged
        assert "test_flag" not in SERVER_FEATURE_FLAGS

    def test_get_server_features_contains_expected_flags(self):
        """Test that server features contain expected flags."""
        features = get_server_features()
        assert "supports_preview_metadata" in features
        assert features["supports_preview_metadata"] is True
        assert "supports_model_type_tags" in features
        assert features["supports_model_type_tags"] is True
        assert "max_upload_size" in features
        assert isinstance(features["max_upload_size"], (int, float))

    def test_get_connection_feature_with_missing_sid(self):
        """Test getting feature for non-existent session ID."""
        sockets_metadata = {}
        result = get_connection_feature(sockets_metadata, "missing_sid", "some_feature")
        assert result is False  # Default value

    def test_get_connection_feature_with_custom_default(self):
        """Test getting feature with custom default value."""
        sockets_metadata = {}
        result = get_connection_feature(
            sockets_metadata, "missing_sid", "some_feature", default="custom_default"
        )
        assert result == "custom_default"

    def test_get_connection_feature_with_feature_flags(self):
        """Test getting feature from connection with feature flags."""
        sockets_metadata = {
            "sid1": {
                "feature_flags": {
                    "supports_preview_metadata": True,
                    "custom_feature": "value",
                },
            }
        }
        result = get_connection_feature(sockets_metadata, "sid1", "supports_preview_metadata")
        assert result is True

        result = get_connection_feature(sockets_metadata, "sid1", "custom_feature")
        assert result == "value"

    def test_get_connection_feature_missing_feature(self):
        """Test getting non-existent feature from connection."""
        sockets_metadata = {
            "sid1": {"feature_flags": {"existing_feature": True}}
        }
        result = get_connection_feature(sockets_metadata, "sid1", "missing_feature")
        assert result is False

    def test_supports_feature_returns_boolean(self):
        """Test that supports_feature always returns boolean."""
        sockets_metadata = {
            "sid1": {
                "feature_flags": {
                    "bool_feature": True,
                    "string_feature": "value",
                    "none_feature": None,
                },
            }
        }

        # True boolean feature
        assert supports_feature(sockets_metadata, "sid1", "bool_feature") is True

        # Non-boolean values should return False
        assert supports_feature(sockets_metadata, "sid1", "string_feature") is False
        assert supports_feature(sockets_metadata, "sid1", "none_feature") is False
        assert supports_feature(sockets_metadata, "sid1", "missing_feature") is False

    def test_supports_feature_with_missing_connection(self):
        """Test supports_feature with missing connection."""
        sockets_metadata = {}
        assert supports_feature(sockets_metadata, "missing_sid", "any_feature") is False

    def test_empty_feature_flags_dict(self):
        """Test connection with empty feature flags dictionary."""
        sockets_metadata = {"sid1": {"feature_flags": {}}}
        result = get_connection_feature(sockets_metadata, "sid1", "any_feature")
        assert result is False
        assert supports_feature(sockets_metadata, "sid1", "any_feature") is False


class TestCoerceFlagValue:
    """Test suite for _coerce_flag_value."""

    def test_registered_bool_true(self):
        assert _coerce_flag_value("show_signin_button", "true") is True
        assert _coerce_flag_value("show_signin_button", "True") is True

    def test_registered_bool_false(self):
        assert _coerce_flag_value("show_signin_button", "false") is False
        assert _coerce_flag_value("show_signin_button", "FALSE") is False

    def test_unregistered_key_stays_string(self):
        assert _coerce_flag_value("unknown_flag", "true") == "true"
        assert _coerce_flag_value("unknown_flag", "42") == "42"

    def test_bool_typo_raises(self):
        """Strict bool: typos like 'ture' or 'yes' must raise so the flag can be dropped."""
        with pytest.raises(ValueError):
            _coerce_flag_value("show_signin_button", "ture")
        with pytest.raises(ValueError):
            _coerce_flag_value("show_signin_button", "yes")
        with pytest.raises(ValueError):
            _coerce_flag_value("show_signin_button", "1")
        with pytest.raises(ValueError):
            _coerce_flag_value("show_signin_button", "")

    def test_failed_int_coercion_raises(self, monkeypatch):
        """Malformed values for typed flags must raise; caller decides what to do."""
        monkeypatch.setitem(
            CLI_FEATURE_FLAG_REGISTRY,
            "test_int_flag",
            {"type": "int", "default": 0, "description": "test"},
        )
        with pytest.raises(ValueError):
            _coerce_flag_value("test_int_flag", "not_a_number")


class TestParseCliFeatureFlags:
    """Test suite for _parse_cli_feature_flags."""

    def test_single_flag(self, monkeypatch):
        monkeypatch.setattr("comfy_api.feature_flags.args", type("Args", (), {"feature_flag": ["show_signin_button=true"]})())
        result = _parse_cli_feature_flags()
        assert result == {"show_signin_button": True}

    def test_missing_equals_defaults_to_true(self, monkeypatch):
        """Bare flag without '=' is treated as the string 'true' (and coerced if registered)."""
        monkeypatch.setattr("comfy_api.feature_flags.args", type("Args", (), {"feature_flag": ["show_signin_button", "valid=1"]})())
        result = _parse_cli_feature_flags()
        assert result == {"show_signin_button": True, "valid": "1"}

    def test_empty_key_skipped(self, monkeypatch):
        monkeypatch.setattr("comfy_api.feature_flags.args", type("Args", (), {"feature_flag": ["=value", "valid=1"]})())
        result = _parse_cli_feature_flags()
        assert result == {"valid": "1"}

    def test_invalid_bool_value_dropped(self, monkeypatch, caplog):
        """A typo'd bool value must be dropped entirely, not silently set to False
        and not stored as a raw string. A warning must be logged."""
        monkeypatch.setattr(
            "comfy_api.feature_flags.args",
            type("Args", (), {"feature_flag": ["show_signin_button=ture", "valid=1"]})(),
        )
        with caplog.at_level("WARNING"):
            result = _parse_cli_feature_flags()
        assert result == {"valid": "1"}
        assert "show_signin_button" not in result
        assert any("show_signin_button" in r.message and "drop" in r.message.lower() for r in caplog.records)


class TestCliFeatureFlagRegistry:
    """Test suite for the CLI feature flag registry."""

    def test_registry_entries_have_required_fields(self):
        for key, info in CLI_FEATURE_FLAG_REGISTRY.items():
            assert "type" in info, f"{key} missing 'type'"
            assert "default" in info, f"{key} missing 'default'"
            assert "description" in info, f"{key} missing 'description'"


class TestComfyApiEnv:
    """--comfy-api-base staging-tier detection + testenv main-host -> -registry rewrite."""

    @pytest.mark.parametrize(
        "url, expected",
        [
            # testenv friendly main host -> comfy-api -registry sibling (slash trimmed)
            ("https://pr-4398.testenvs.comfy.org", "https://pr-4398-registry.testenvs.comfy.org"),
            ("https://pr-4398.testenvs.comfy.org/", "https://pr-4398-registry.testenvs.comfy.org"),
            ("https://pr-4398-registry.testenvs.comfy.org", "https://pr-4398-registry.testenvs.comfy.org"),
            # staging + everything else -> unchanged (no -registry split)
            ("https://stagingapi.comfy.org", "https://stagingapi.comfy.org"),
            ("https://api.comfy.org", "https://api.comfy.org"),
            ("https://pr-1.testenvs.comfy.org.evil.com", "https://pr-1.testenvs.comfy.org.evil.com"),
            ("", ""),
        ],
    )
    def test_normalize_comfy_api_base(self, url, expected):
        assert normalize_comfy_api_base(url) == expected

    def test_config_for_staging_tier_else_none(self):
        # ephemeral testenv: friendly main host -> -registry, staging platform, dev Firebase env
        eph = environment_overrides_for_base("https://pr-1234.testenvs.comfy.org/")
        assert eph["comfy_api_base_url"] == "https://pr-1234-registry.testenvs.comfy.org"
        assert eph["comfy_platform_base_url"] == "https://stagingplatform.comfy.org"
        assert eph["firebase_env"] == "dev"
        # staging api host: emitted as-is
        stg = environment_overrides_for_base("https://stagingapi.comfy.org")
        assert stg["comfy_api_base_url"] == "https://stagingapi.comfy.org"
        assert stg["comfy_platform_base_url"] == "https://stagingplatform.comfy.org"
        assert stg["firebase_env"] == "dev"
        # prod / unknown: nothing
        assert environment_overrides_for_base("https://api.comfy.org") is None

    def test_environment_overrides_only_for_staging_tier(self, monkeypatch):
        def set_base(url):
            monkeypatch.setattr(
                "comfy.comfy_api_env.args",
                type("Args", (), {"comfy_api_base": url})(),
            )

        # The overrides merged into the HTTP /features response are present for staging-tier bases...
        set_base("https://stagingapi.comfy.org")
        assert "comfy_api_base_url" in get_environment_overrides()
        set_base("https://pr-7.testenvs.comfy.org")
        assert "comfy_api_base_url" in get_environment_overrides()
        # ...but never for prod.
        set_base("https://api.comfy.org")
        assert get_environment_overrides() is None

    def test_server_features_never_carry_env_overrides(self, monkeypatch):
        """The WebSocket capability handshake must stay free of routing keys."""
        monkeypatch.setattr(
            "comfy.comfy_api_env.args",
            type("Args", (), {"comfy_api_base": "https://pr-7.testenvs.comfy.org"})(),
        )
        features = get_server_features()
        assert "comfy_api_base_url" not in features
        assert "comfy_platform_base_url" not in features
        assert "firebase_env" not in features