from __future__ import annotations

from comfy_engine.scripts.gateway_auth import (
    extract_gateway_api_key,
    is_gateway_request_authorized,
    load_gateway_api_key,
)


def test_load_gateway_api_key_prefers_explicit_gateway_env() -> None:
    env = {
        "DREAM_WEAVER_API_KEY": "fallback-key",
        "DREAM_WEAVER_GATEWAY_API_KEY": "primary-key",
    }
    assert load_gateway_api_key(env) == "primary-key"


def test_extract_gateway_api_key_supports_dedicated_headers() -> None:
    assert extract_gateway_api_key({"x-dream-weaver-api-key": "dw-key"}) == "dw-key"
    assert extract_gateway_api_key({"x-api-key": "legacy-key"}) == "legacy-key"


def test_extract_gateway_api_key_supports_bearer_authorization() -> None:
    assert extract_gateway_api_key({"authorization": "Bearer shared-key"}) == "shared-key"


def test_gateway_auth_allows_open_gateways_and_blocks_wrong_keys() -> None:
    assert is_gateway_request_authorized({}, None) is True
    assert is_gateway_request_authorized({"x-dream-weaver-api-key": "correct"}, "correct") is True
    assert is_gateway_request_authorized({"authorization": "Bearer wrong"}, "correct") is False