114 lines
3.8 KiB
Python
114 lines
3.8 KiB
Python
"""Tests for the generic remote worker framing protocol."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import io
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
|
|
def test_binary_frame_round_trip_preserves_arbitrary_bytes(remote_protocol_module: Any) -> None:
|
|
"""Binary tensor/result payloads must not pass through base64 or text."""
|
|
payload = bytes(range(256)) * 1024
|
|
encoded = remote_protocol_module.encode_frame(
|
|
remote_protocol_module.RemoteFrameKind.RESULT,
|
|
payload,
|
|
)
|
|
|
|
decoded = remote_protocol_module.read_frame(io.BytesIO(encoded))
|
|
|
|
assert decoded == (remote_protocol_module.RemoteFrameKind.RESULT, payload)
|
|
assert len(encoded) == len(payload) + 17
|
|
|
|
|
|
def test_json_frames_use_canonical_object_payloads(remote_protocol_module: Any) -> None:
|
|
"""Control frames should round-trip deterministic JSON objects."""
|
|
encoded = remote_protocol_module.encode_json_frame(
|
|
remote_protocol_module.RemoteFrameKind.CANCEL,
|
|
{"invocation_id": "RIV_test", "reason": "user"},
|
|
)
|
|
kind, payload = remote_protocol_module.read_frame(io.BytesIO(encoded))
|
|
|
|
assert kind is remote_protocol_module.RemoteFrameKind.CANCEL
|
|
assert remote_protocol_module.decode_json_payload(payload) == {
|
|
"invocation_id": "RIV_test",
|
|
"reason": "user",
|
|
}
|
|
|
|
|
|
def test_truncated_frame_is_rejected(remote_protocol_module: Any) -> None:
|
|
"""A transport disconnect must not be mistaken for a valid empty result."""
|
|
encoded = remote_protocol_module.encode_frame(
|
|
remote_protocol_module.RemoteFrameKind.INPUTS,
|
|
b"abcdef",
|
|
)
|
|
|
|
with pytest.raises(remote_protocol_module.RemoteProtocolError, match="outstanding"):
|
|
remote_protocol_module.read_frame(io.BytesIO(encoded[:-2]))
|
|
|
|
|
|
def test_frame_io_reports_incremental_payload_bytes(remote_protocol_module: Any) -> None:
|
|
"""Large framed transfers should expose byte progress in bounded chunks."""
|
|
payload = b"x" * (2 * 1024 * 1024 + 17)
|
|
encoded_stream = io.BytesIO()
|
|
written: list[tuple[Any, int, int]] = []
|
|
remote_protocol_module.write_frame(
|
|
encoded_stream,
|
|
remote_protocol_module.RemoteFrameKind.RESULT,
|
|
payload,
|
|
progress_callback=lambda kind, current, total: written.append(
|
|
(kind, current, total)
|
|
),
|
|
)
|
|
read: list[tuple[Any, int, int]] = []
|
|
|
|
decoded = remote_protocol_module.read_frame(
|
|
io.BytesIO(encoded_stream.getvalue()),
|
|
progress_callback=lambda kind, current, total: read.append(
|
|
(kind, current, total)
|
|
),
|
|
)
|
|
|
|
assert decoded == (remote_protocol_module.RemoteFrameKind.RESULT, payload)
|
|
assert [current for _kind, current, _total in written] == [
|
|
1024 * 1024,
|
|
2 * 1024 * 1024,
|
|
len(payload),
|
|
]
|
|
assert read == written
|
|
|
|
|
|
def test_relay_frame_streams_without_reassembling_payload(
|
|
remote_protocol_module: Any,
|
|
) -> None:
|
|
"""The SSH relay should forward large frames a chunk at a time."""
|
|
payload = b"z" * (2 * 1024 * 1024 + 7)
|
|
encoded = remote_protocol_module.encode_frame(
|
|
remote_protocol_module.RemoteFrameKind.INPUTS,
|
|
payload,
|
|
)
|
|
forwarded_chunks: list[bytes] = []
|
|
|
|
kind = remote_protocol_module.relay_frame(
|
|
io.BytesIO(encoded),
|
|
lambda chunk: forwarded_chunks.append(bytes(chunk)),
|
|
)
|
|
|
|
assert kind is remote_protocol_module.RemoteFrameKind.INPUTS
|
|
assert b"".join(forwarded_chunks) == encoded
|
|
assert len(forwarded_chunks) == 4
|
|
|
|
|
|
def test_protocol_exports_shared_prompt_boundary_constants(
|
|
remote_protocol_module: Any,
|
|
) -> None:
|
|
"""Local and cloud runtimes must share exact prompt payload wire keys."""
|
|
assert (
|
|
remote_protocol_module.BOUNDARY_INPUT_SIGNATURES_KEY
|
|
== "__comfy_modal_boundary_input_signatures__"
|
|
)
|
|
assert remote_protocol_module.PRIMITIVE_WIDGET_INPUT_TYPES == frozenset(
|
|
{"INT", "FLOAT", "BOOLEAN", "STRING"}
|
|
)
|