774 lines
25 KiB
Python
774 lines
25 KiB
Python
"""Tests for the Modal hosted-model endpoint node."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import subprocess
|
|
from dataclasses import dataclass
|
|
from typing import Any, Self
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
|
|
class FakeKeyringError(Exception):
|
|
"""Error raised by the fake keyring backend."""
|
|
|
|
|
|
class FakeKeyring:
|
|
"""Small in-memory implementation of the keyring API used by the node."""
|
|
|
|
errors = type("Errors", (), {"KeyringError": FakeKeyringError})
|
|
|
|
def __init__(self) -> None:
|
|
"""Create an empty fake credential vault."""
|
|
self.values: dict[tuple[str, str], str] = {}
|
|
|
|
def get_password(self, service: str, username: str) -> str | None:
|
|
"""Return one fake secret value."""
|
|
return self.values.get((service, username))
|
|
|
|
def set_password(self, service: str, username: str, value: str) -> None:
|
|
"""Store one fake secret value."""
|
|
self.values[(service, username)] = value
|
|
|
|
def delete_password(self, service: str, username: str) -> None:
|
|
"""Delete one fake secret value."""
|
|
self.values.pop((service, username), None)
|
|
|
|
|
|
@dataclass
|
|
class FakeInputFile:
|
|
"""Pydantic-like file input used by the built-in OpenAI file node."""
|
|
|
|
filename: str
|
|
file_data: str
|
|
|
|
def model_dump(self, exclude_none: bool = True) -> dict[str, str]:
|
|
"""Return the fields exposed by ComfyUI's InputFileContent model."""
|
|
del exclude_none
|
|
return {
|
|
"type": "input_file",
|
|
"filename": self.filename,
|
|
"file_data": self.file_data,
|
|
}
|
|
|
|
|
|
class FakeResponseContent:
|
|
"""Async byte stream for a fake aiohttp response."""
|
|
|
|
def __init__(self, payload: dict[str, Any] | bytes) -> None:
|
|
"""Serialize one response object."""
|
|
self._body = (
|
|
payload
|
|
if isinstance(payload, bytes)
|
|
else json.dumps(payload).encode("utf-8")
|
|
)
|
|
|
|
async def iter_chunked(self, chunk_size: int) -> Any:
|
|
"""Yield the serialized body in bounded pieces."""
|
|
for offset in range(0, len(self._body), chunk_size):
|
|
yield self._body[offset : offset + chunk_size]
|
|
|
|
|
|
class FakeResponse:
|
|
"""Async context manager matching the aiohttp response surface used by the client."""
|
|
|
|
def __init__(
|
|
self, payload: dict[str, Any] | bytes, status: int = 200
|
|
) -> None:
|
|
"""Configure the response JSON and status code."""
|
|
self.status = status
|
|
self.content = FakeResponseContent(payload)
|
|
|
|
async def __aenter__(self) -> Self:
|
|
"""Enter the fake response context."""
|
|
return self
|
|
|
|
async def __aexit__(self, *exc_info: object) -> None:
|
|
"""Leave the fake response context."""
|
|
del exc_info
|
|
|
|
|
|
class FakeSession:
|
|
"""Capture one aiohttp-style request call."""
|
|
|
|
def __init__(self, response: FakeResponse | list[FakeResponse]) -> None:
|
|
"""Configure one or more responses returned by request calls."""
|
|
self.responses = response if isinstance(response, list) else [response]
|
|
self.calls: list[dict[str, Any]] = []
|
|
|
|
def request(self, method: str, url: str, **kwargs: Any) -> FakeResponse:
|
|
"""Record and return one request context manager."""
|
|
self.calls.append({"method": method, "url": url, **kwargs})
|
|
return self.responses.pop(0)
|
|
|
|
|
|
def _credentials(modal_endpoint_module: Any) -> Any:
|
|
"""Return a syntactically valid test proxy-token pair."""
|
|
return modal_endpoint_module.ModalProxyCredentials(
|
|
key="wk-test-key",
|
|
secret="ws-test-secret",
|
|
)
|
|
|
|
|
|
def test_endpoint_url_normalization_accepts_modal_direct_suffixes(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Accept the endpoint forms users are likely to paste from Modal."""
|
|
endpoint = "https://example--model.us-west.modal.direct"
|
|
|
|
assert modal_endpoint_module._normalize_endpoint_url(endpoint) == endpoint
|
|
assert modal_endpoint_module._normalize_endpoint_url(f"{endpoint}/v1") == endpoint
|
|
assert (
|
|
modal_endpoint_module._normalize_endpoint_url(f"{endpoint}/v1/chat/completions")
|
|
== endpoint
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"endpoint",
|
|
[
|
|
"http://example--model.us-west.modal.direct",
|
|
"https://example.com",
|
|
"https://example--model.us-west.modal.direct@attacker.example",
|
|
"https://example--model.us-west.modal.direct/other",
|
|
"https://example--model.us-west.modal.direct?next=https://attacker.example",
|
|
],
|
|
)
|
|
def test_endpoint_url_normalization_rejects_credential_exfiltration_targets(
|
|
modal_endpoint_module: Any,
|
|
endpoint: str,
|
|
) -> None:
|
|
"""Never attach Modal credentials to a noncanonical destination."""
|
|
with pytest.raises(ValueError):
|
|
modal_endpoint_module._normalize_endpoint_url(endpoint)
|
|
|
|
|
|
def test_secret_manager_round_trips_credentials(modal_endpoint_module: Any) -> None:
|
|
"""Persist both proxy-token values under the ComfyUI-specific vault service."""
|
|
keyring = FakeKeyring()
|
|
manager = modal_endpoint_module.ComfyUISecretManager(keyring)
|
|
credentials = _credentials(modal_endpoint_module)
|
|
|
|
assert manager.load() is None
|
|
manager.ensure_writable()
|
|
assert keyring.values == {}
|
|
manager.save(credentials)
|
|
|
|
assert manager.load() == credentials
|
|
assert set(keyring.values.values()) == {credentials.key, credentials.secret}
|
|
|
|
|
|
def test_secret_manager_rejects_partial_pair(modal_endpoint_module: Any) -> None:
|
|
"""Do not create a replacement token over a damaged stored pair."""
|
|
keyring = FakeKeyring()
|
|
manager = modal_endpoint_module.ComfyUISecretManager(keyring)
|
|
keyring.set_password("ComfyUI Modal-Sync", "MODAL_KEY", "wk-only")
|
|
|
|
with pytest.raises(RuntimeError, match="only one half"):
|
|
manager.load()
|
|
|
|
|
|
def test_credential_resolver_prefers_environment(
|
|
modal_endpoint_module: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Use explicit environment credentials without reading or writing the vault."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
monkeypatch.setenv("MODAL_KEY", credentials.key)
|
|
monkeypatch.setenv("MODAL_SECRET", credentials.secret)
|
|
|
|
class Unused:
|
|
"""Fail when an environment-only resolution touches a fallback."""
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
"""Reject every fallback operation."""
|
|
raise AssertionError(f"unexpected fallback access: {name}")
|
|
|
|
resolver = modal_endpoint_module.ModalCredentialResolver(Unused(), Unused())
|
|
|
|
assert resolver.resolve() == credentials
|
|
|
|
|
|
def test_credential_resolver_creates_and_stores_missing_pair(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Create exactly one pair when neither environment nor vault has credentials."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
|
|
class Store:
|
|
"""Record credential-store interactions."""
|
|
|
|
saved: Any = None
|
|
|
|
def load(self) -> None:
|
|
"""Report an empty vault."""
|
|
|
|
def save(self, value: Any) -> None:
|
|
"""Record the newly created credential pair."""
|
|
self.saved = value
|
|
|
|
def ensure_writable(self) -> None:
|
|
"""Confirm the fake store can accept the new token."""
|
|
|
|
class Creator:
|
|
"""Return one deterministic proxy token."""
|
|
|
|
calls = 0
|
|
|
|
def create(self) -> Any:
|
|
"""Create one test pair."""
|
|
self.calls += 1
|
|
return credentials
|
|
|
|
store = Store()
|
|
creator = Creator()
|
|
|
|
assert (
|
|
modal_endpoint_module.ModalCredentialResolver(store, creator).resolve()
|
|
== credentials
|
|
)
|
|
assert creator.calls == 1
|
|
assert store.saved == credentials
|
|
|
|
|
|
def test_credential_resolver_authorizes_vault_credentials(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Authorize node-managed scoped tokens for the endpoint environment."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
|
|
class Store:
|
|
"""Return a credential pair previously created by the node."""
|
|
|
|
def load(self) -> Any:
|
|
"""Return stored credentials."""
|
|
return credentials
|
|
|
|
class Authorizer:
|
|
"""Record the environment association request."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create an empty authorization call log."""
|
|
self.calls: list[tuple[str, str]] = []
|
|
|
|
def allow(self, token_key: str, environment: str) -> None:
|
|
"""Record one authorization operation."""
|
|
self.calls.append((token_key, environment))
|
|
|
|
authorizer = Authorizer()
|
|
resolver = modal_endpoint_module.ModalCredentialResolver(
|
|
Store(),
|
|
object(),
|
|
authorizer=authorizer,
|
|
environment="ComfyUI",
|
|
)
|
|
|
|
assert resolver.resolve() == credentials
|
|
assert authorizer.calls == [(credentials.key, "ComfyUI")]
|
|
|
|
|
|
def test_credential_resolver_treats_blank_environment_as_main(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Keep workflows saved with an empty advanced environment compatible."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
|
|
class Store:
|
|
"""Return a credential pair previously created by the node."""
|
|
|
|
def load(self) -> Any:
|
|
"""Return stored credentials."""
|
|
return credentials
|
|
|
|
class Authorizer:
|
|
"""Record the normalized environment association."""
|
|
|
|
environment: str | None = None
|
|
|
|
def allow(self, token_key: str, environment: str) -> None:
|
|
"""Record one authorization operation."""
|
|
assert token_key == credentials.key
|
|
self.environment = environment
|
|
|
|
authorizer = Authorizer()
|
|
resolver = modal_endpoint_module.ModalCredentialResolver(
|
|
Store(), object(), authorizer=authorizer, environment=" "
|
|
)
|
|
|
|
assert resolver.resolve() == credentials
|
|
assert authorizer.environment == "main"
|
|
|
|
|
|
def test_credential_resolver_saves_new_token_before_authorization(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Keep the one-time secret recoverable when environment authorization fails."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
events: list[str] = []
|
|
|
|
class Store:
|
|
"""Record persistence ordering for a new token."""
|
|
|
|
def load(self) -> None:
|
|
"""Report no stored credentials."""
|
|
|
|
def ensure_writable(self) -> None:
|
|
"""Accept the credential-vault probe."""
|
|
|
|
def save(self, value: Any) -> None:
|
|
"""Record persistence of the new token."""
|
|
assert value == credentials
|
|
events.append("saved")
|
|
|
|
class Creator:
|
|
"""Return one deterministic one-time token pair."""
|
|
|
|
def create(self) -> Any:
|
|
"""Return generated credentials."""
|
|
return credentials
|
|
|
|
class Authorizer:
|
|
"""Fail after verifying the secret was persisted."""
|
|
|
|
def allow(self, token_key: str, environment: str) -> None:
|
|
"""Simulate an RBAC association failure."""
|
|
assert token_key == credentials.key
|
|
assert environment == "main"
|
|
assert events == ["saved"]
|
|
raise RuntimeError("authorization failed")
|
|
|
|
resolver = modal_endpoint_module.ModalCredentialResolver(
|
|
Store(), Creator(), authorizer=Authorizer()
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="authorization failed"):
|
|
resolver.resolve()
|
|
assert events == ["saved"]
|
|
|
|
|
|
def test_credential_resolver_checks_vault_before_creating_token(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Do not mint a one-time proxy token when secure persistence is unavailable."""
|
|
|
|
class Store:
|
|
"""Represent a readable but non-writable credential vault."""
|
|
|
|
def load(self) -> None:
|
|
"""Report no stored pair."""
|
|
|
|
def ensure_writable(self) -> None:
|
|
"""Reject writes before token creation begins."""
|
|
raise RuntimeError("vault unavailable")
|
|
|
|
class Creator:
|
|
"""Fail if token creation is attempted."""
|
|
|
|
def create(self) -> Any:
|
|
"""Reject an unsafe token-creation attempt."""
|
|
raise AssertionError("token creation should not run")
|
|
|
|
with pytest.raises(RuntimeError, match="vault unavailable"):
|
|
modal_endpoint_module.ModalCredentialResolver(Store(), Creator()).resolve()
|
|
|
|
|
|
def test_cli_creator_parses_json_response(modal_endpoint_module: Any) -> None:
|
|
"""Parse Modal's documented --json response without exposing values elsewhere."""
|
|
output = json.dumps({"Modal-Key": "wk-created", "Modal-Secret": "ws-created"})
|
|
|
|
assert modal_endpoint_module.ModalCliProxyTokenCreator._parse_credentials(
|
|
output
|
|
) == (modal_endpoint_module.ModalProxyCredentials("wk-created", "ws-created"))
|
|
|
|
|
|
def test_cli_creator_falls_back_from_outdated_modal(
|
|
modal_endpoint_module: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Use uvx when an installed Modal CLI predates workspace proxy-token commands."""
|
|
creator = modal_endpoint_module.ModalCliProxyTokenCreator()
|
|
commands = [["old-modal"], ["uvx-modal"]]
|
|
results = iter(
|
|
[
|
|
subprocess.CompletedProcess(
|
|
commands[0],
|
|
2,
|
|
stdout="",
|
|
stderr="No such command 'workspace'.",
|
|
),
|
|
subprocess.CompletedProcess(
|
|
commands[1],
|
|
0,
|
|
stdout=json.dumps({"Modal-Key": "wk-new", "Modal-Secret": "ws-new"}),
|
|
stderr="",
|
|
),
|
|
]
|
|
)
|
|
monkeypatch.setattr(creator, "_candidate_commands", lambda: commands)
|
|
monkeypatch.setattr(creator, "_run_command", lambda command: next(results))
|
|
|
|
assert creator.create() == modal_endpoint_module.ModalProxyCredentials(
|
|
"wk-new", "ws-new"
|
|
)
|
|
|
|
|
|
def test_cli_authorizer_falls_back_and_caches_success(
|
|
modal_endpoint_module: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Use a current CLI once for a scoped token and avoid per-request CLI latency."""
|
|
authorizer = modal_endpoint_module.ModalCliProxyTokenAuthorizer()
|
|
commands = [["old-modal"], ["uvx-modal"]]
|
|
results = iter(
|
|
[
|
|
subprocess.CompletedProcess(
|
|
commands[0],
|
|
2,
|
|
stdout="",
|
|
stderr="No such command 'workspace'.",
|
|
),
|
|
subprocess.CompletedProcess(commands[1], 0, stdout="allowed", stderr=""),
|
|
]
|
|
)
|
|
calls: list[list[str]] = []
|
|
monkeypatch.setattr(
|
|
authorizer,
|
|
"_candidate_commands",
|
|
lambda token_key, environment: commands,
|
|
)
|
|
|
|
def run_command(command: list[str]) -> subprocess.CompletedProcess[str]:
|
|
"""Record one fake CLI invocation."""
|
|
calls.append(command)
|
|
return next(results)
|
|
|
|
monkeypatch.setattr(
|
|
modal_endpoint_module.ModalCliProxyTokenCreator,
|
|
"_run_command",
|
|
run_command,
|
|
)
|
|
|
|
authorizer.allow("wk-authorizer-test", "main")
|
|
authorizer.allow("wk-authorizer-test", "main")
|
|
|
|
assert calls == commands
|
|
|
|
|
|
def test_image_batch_becomes_inline_png_content_parts(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Encode every BHWC tensor item as an OpenAI-compatible image_url block."""
|
|
images = torch.zeros((2, 3, 4, 3), dtype=torch.float32)
|
|
|
|
content_parts = modal_endpoint_module._image_content_parts(images)
|
|
|
|
assert len(content_parts) == 2
|
|
data_uri = content_parts[0]["image_url"]["url"]
|
|
assert data_uri.startswith("data:image/png;base64,")
|
|
assert base64.b64decode(data_uri.split(",", 1)[1]).startswith(b"\x89PNG\r\n\x1a\n")
|
|
|
|
|
|
def test_builtin_openai_file_input_becomes_chat_file_part(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Accept the existing OpenAI ChatGPT Input Files node output directly."""
|
|
file_input = FakeInputFile(
|
|
filename="notes.txt",
|
|
file_data="data:text/plain;base64,SGVsbG8=",
|
|
)
|
|
|
|
assert modal_endpoint_module._file_content_part(file_input) == {
|
|
"type": "file",
|
|
"file": {
|
|
"filename": "notes.txt",
|
|
"file_data": "data:text/plain;base64,SGVsbG8=",
|
|
},
|
|
}
|
|
|
|
|
|
def test_request_payload_matches_chat_completions_contract(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Build system and multimodal user messages with explicit generation options."""
|
|
payload = modal_endpoint_module.ModalEndpointClient._request_payload(
|
|
prompt="Describe this",
|
|
model="org/model",
|
|
images=torch.zeros((1, 2, 2, 3)),
|
|
files=[FakeInputFile("notes.txt", "data:text/plain;base64,SGVsbG8=")],
|
|
system_prompt="Be concise",
|
|
max_tokens=123,
|
|
temperature=0.25,
|
|
)
|
|
|
|
assert payload["model"] == "org/model"
|
|
assert payload["stream"] is False
|
|
assert payload["max_tokens"] == 123
|
|
assert payload["temperature"] == 0.25
|
|
assert payload["messages"][0] == {"role": "system", "content": "Be concise"}
|
|
assert [part["type"] for part in payload["messages"][1]["content"]] == [
|
|
"text",
|
|
"image_url",
|
|
"file",
|
|
]
|
|
|
|
|
|
def test_http_client_authenticates_without_following_redirects(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Attach both proxy-token headers only to the validated canonical endpoint URL."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
client = modal_endpoint_module.ModalEndpointClient(
|
|
"https://example--model.us-west.modal.direct/v1",
|
|
credentials,
|
|
timeout_seconds=30,
|
|
)
|
|
session = FakeSession(FakeResponse({"data": []}))
|
|
|
|
response = asyncio.run(client._request_json(session, "GET", "/v1/models"))
|
|
|
|
assert response == {"data": []}
|
|
assert session.calls == [
|
|
{
|
|
"method": "GET",
|
|
"url": "https://example--model.us-west.modal.direct/v1/models",
|
|
"headers": {
|
|
"Content-Type": "application/json",
|
|
"Modal-Key": credentials.key,
|
|
"Modal-Secret": credentials.secret,
|
|
},
|
|
"json": None,
|
|
"allow_redirects": False,
|
|
}
|
|
]
|
|
|
|
|
|
def test_endpoint_error_message_redacts_modal_proxy_tokens(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Do not repeat credential identifiers from provider errors into ComfyUI logs."""
|
|
detail = modal_endpoint_module._response_error_message(
|
|
{"error": "Webhook token not found: wk-sensitive-token"}
|
|
)
|
|
|
|
assert detail == '{"error": "Webhook token not found: [redacted]"}'
|
|
|
|
|
|
def test_http_client_retries_empty_503_until_modal_replica_is_ready(
|
|
modal_endpoint_module: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Treat Modal Server's empty 503 as a cold-start signal, not malformed JSON."""
|
|
client = modal_endpoint_module.ModalEndpointClient(
|
|
"https://example--model.us-west.modal.direct",
|
|
_credentials(modal_endpoint_module),
|
|
timeout_seconds=30,
|
|
)
|
|
session = FakeSession(
|
|
[
|
|
FakeResponse(b"", status=503),
|
|
FakeResponse({"data": [{"id": "org/model"}]}),
|
|
]
|
|
)
|
|
delays: list[float] = []
|
|
|
|
async def record_sleep(delay: float) -> None:
|
|
"""Record retry backoff without delaying the test."""
|
|
delays.append(delay)
|
|
|
|
monkeypatch.setattr(modal_endpoint_module.asyncio, "sleep", record_sleep)
|
|
|
|
response = asyncio.run(client._request_json(session, "GET", "/v1/models"))
|
|
|
|
assert response == {"data": [{"id": "org/model"}]}
|
|
assert delays == [1.0]
|
|
assert len(session.calls) == 2
|
|
|
|
|
|
def test_http_client_reports_status_for_empty_nonretryable_response(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Preserve the HTTP status when an endpoint error has no JSON body."""
|
|
client = modal_endpoint_module.ModalEndpointClient(
|
|
"https://example--model.us-west.modal.direct",
|
|
_credentials(modal_endpoint_module),
|
|
timeout_seconds=30,
|
|
)
|
|
session = FakeSession(FakeResponse(b"", status=502))
|
|
|
|
with pytest.raises(RuntimeError, match="HTTP 502 with an empty response body"):
|
|
asyncio.run(client._request_json(session, "GET", "/v1/models"))
|
|
|
|
|
|
def test_http_client_does_not_retry_application_json_503(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Surface an application-level 503 body instead of treating it as scale-up."""
|
|
client = modal_endpoint_module.ModalEndpointClient(
|
|
"https://example--model.us-west.modal.direct",
|
|
_credentials(modal_endpoint_module),
|
|
timeout_seconds=30,
|
|
)
|
|
session = FakeSession(
|
|
FakeResponse({"error": {"message": "model failed"}}, status=503)
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="HTTP 503: model failed"):
|
|
asyncio.run(client._request_json(session, "GET", "/v1/models"))
|
|
assert len(session.calls) == 1
|
|
|
|
|
|
def test_http_client_bounds_503_retries_by_total_timeout(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Stop cold-start polling at the node's total request deadline."""
|
|
client = modal_endpoint_module.ModalEndpointClient(
|
|
"https://example--model.us-west.modal.direct",
|
|
_credentials(modal_endpoint_module),
|
|
timeout_seconds=0.01,
|
|
)
|
|
session = FakeSession(FakeResponse(b"", status=503))
|
|
|
|
with pytest.raises(TimeoutError, match="stayed unavailable with HTTP 503"):
|
|
asyncio.run(client._request_json(session, "GET", "/v1/models"))
|
|
|
|
|
|
def test_http_client_reports_bounded_non_json_error_body(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Include safe upstream text when an endpoint does not return JSON."""
|
|
client = modal_endpoint_module.ModalEndpointClient(
|
|
"https://example--model.us-west.modal.direct",
|
|
_credentials(modal_endpoint_module),
|
|
timeout_seconds=30,
|
|
)
|
|
session = FakeSession(FakeResponse(b"upstream unavailable", status=502))
|
|
|
|
with pytest.raises(
|
|
RuntimeError,
|
|
match="HTTP 502 with a non-JSON response: upstream unavailable",
|
|
):
|
|
asyncio.run(client._request_json(session, "GET", "/v1/models"))
|
|
|
|
|
|
def test_chat_response_text_supports_string_and_part_lists(
|
|
modal_endpoint_module: Any,
|
|
) -> None:
|
|
"""Extract common OpenAI-compatible assistant content representations."""
|
|
assert (
|
|
modal_endpoint_module._chat_response_text(
|
|
{"choices": [{"message": {"content": "hello"}}]}
|
|
)
|
|
== "hello"
|
|
)
|
|
assert (
|
|
modal_endpoint_module._chat_response_text(
|
|
{
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": [
|
|
{"type": "text", "text": "hello "},
|
|
{"type": "text", "text": "world"},
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
)
|
|
== "hello world"
|
|
)
|
|
|
|
|
|
def test_node_schema_matches_builtin_chat_shape(modal_endpoint_module: Any) -> None:
|
|
"""Expose prompt, image, file, endpoint, and model controls as a V3 node."""
|
|
schema = modal_endpoint_module.ModalEndpointChat.define_schema()
|
|
input_ids = [input_spec.id for input_spec in schema.inputs]
|
|
|
|
assert schema.node_id == "ModalEndpointChat"
|
|
assert schema.display_name == "Modal Endpoint Chat"
|
|
assert schema.essentials_category == "Text Generation"
|
|
assert input_ids[:5] == ["prompt", "endpoint_url", "model", "images", "files"]
|
|
assert schema.outputs[0].io_type == "STRING"
|
|
|
|
|
|
def test_extension_exports_endpoint_node(
|
|
extension_package: Any, modal_endpoint_module: Any
|
|
) -> None:
|
|
"""Register the endpoint node through the package's ComfyUI V3 entrypoint module."""
|
|
assert (
|
|
extension_package.ModalEndpointChat is modal_endpoint_module.ModalEndpointChat
|
|
)
|
|
|
|
|
|
def test_node_execute_returns_endpoint_text(
|
|
modal_endpoint_module: Any,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Wire credential resolution and the async client into a standard node output."""
|
|
credentials = _credentials(modal_endpoint_module)
|
|
|
|
class Resolver:
|
|
"""Return deterministic credentials."""
|
|
|
|
def __init__(
|
|
self,
|
|
store: Any,
|
|
creator: Any,
|
|
authorizer: Any,
|
|
environment: str,
|
|
) -> None:
|
|
"""Accept production constructor dependencies."""
|
|
del store, creator
|
|
assert authorizer is not None
|
|
assert environment == "main"
|
|
|
|
def resolve(self) -> Any:
|
|
"""Return the test credential pair."""
|
|
return credentials
|
|
|
|
class Client:
|
|
"""Return a deterministic endpoint completion."""
|
|
|
|
def __init__(
|
|
self, endpoint_url: str, supplied_credentials: Any, timeout: int
|
|
) -> None:
|
|
"""Record the validated constructor arguments."""
|
|
assert endpoint_url.endswith("modal.direct")
|
|
assert supplied_credentials == credentials
|
|
assert timeout == 30
|
|
self._origin = endpoint_url
|
|
|
|
@property
|
|
def endpoint_hostname(self) -> str:
|
|
"""Return the fake endpoint hostname."""
|
|
return "example--model.us-west.modal.direct"
|
|
|
|
async def complete(self, **kwargs: Any) -> str:
|
|
"""Return one fake assistant response."""
|
|
assert kwargs["prompt"] == "hello"
|
|
return "modal reply"
|
|
|
|
monkeypatch.setattr(modal_endpoint_module, "ComfyUISecretManager", lambda: object())
|
|
monkeypatch.setattr(
|
|
modal_endpoint_module, "ModalCliProxyTokenCreator", lambda: object()
|
|
)
|
|
monkeypatch.setattr(
|
|
modal_endpoint_module, "ModalCliProxyTokenAuthorizer", lambda: object()
|
|
)
|
|
monkeypatch.setattr(modal_endpoint_module, "ModalCredentialResolver", Resolver)
|
|
monkeypatch.setattr(modal_endpoint_module, "ModalEndpointClient", Client)
|
|
|
|
output = asyncio.run(
|
|
modal_endpoint_module.ModalEndpointChat.execute(
|
|
prompt="hello",
|
|
endpoint_url="https://example--model.us-west.modal.direct",
|
|
timeout_seconds=30,
|
|
)
|
|
)
|
|
|
|
assert output.result == ("modal reply",)
|