165 lines
5.1 KiB
Python
165 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Protocol schema tests for the streaming server.
|
|
|
|
Covers:
|
|
|
|
* accepted client messages parse into the correct discriminated model
|
|
* unknown ``type`` values raise validation errors
|
|
* server-side messages serialize to the expected wire shape
|
|
* continuation_state field on session_init_v2 carries through
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from fastvideo.entrypoints.streaming.protocol import (
|
|
ContinuationStateSnapshot,
|
|
ErrorMessage,
|
|
GpuAssigned,
|
|
Ltx2SegmentComplete,
|
|
Ltx2SegmentStart,
|
|
Ltx2StreamStart,
|
|
MediaInit,
|
|
MediaSegmentComplete,
|
|
QueueStatus,
|
|
SegmentPromptSource,
|
|
SessionInitV2,
|
|
SnapshotState,
|
|
StepComplete,
|
|
parse_client_message,
|
|
)
|
|
|
|
|
|
class TestClientMessageParsing:
|
|
|
|
def test_session_init_v2_minimal(self):
|
|
parsed = parse_client_message({"type": "session_init_v2"})
|
|
assert isinstance(parsed, SessionInitV2)
|
|
assert parsed.curated_prompts == []
|
|
assert parsed.stream_mode == "av_fmp4"
|
|
|
|
def test_session_init_v2_full(self):
|
|
raw = {
|
|
"type": "session_init_v2",
|
|
"client_id": "client-1",
|
|
"preset": "ltx2_two_stage",
|
|
"preset_label": "2x refine",
|
|
"curated_prompts": ["a fox", "a deer"],
|
|
"enhancement_enabled": True,
|
|
"auto_extension_enabled": False,
|
|
"loop_generation_enabled": False,
|
|
"single_clip_mode": True,
|
|
"stream_mode": "av_fmp4",
|
|
"continuation_state": {
|
|
"kind": "ltx2.v1",
|
|
"payload": {
|
|
"schema_version": 1,
|
|
"segment_index": 2
|
|
},
|
|
},
|
|
}
|
|
parsed = parse_client_message(raw)
|
|
assert isinstance(parsed, SessionInitV2)
|
|
assert parsed.preset == "ltx2_two_stage"
|
|
assert parsed.curated_prompts == ["a fox", "a deer"]
|
|
assert parsed.continuation_state["kind"] == "ltx2.v1"
|
|
|
|
def test_segment_prompt_source(self):
|
|
parsed = parse_client_message({
|
|
"type": "segment_prompt_source",
|
|
"prompt": "hello world",
|
|
"source": "curated",
|
|
"seed": 7,
|
|
})
|
|
assert isinstance(parsed, SegmentPromptSource)
|
|
assert parsed.source == "curated"
|
|
assert parsed.seed == 7
|
|
|
|
def test_snapshot_state(self):
|
|
parsed = parse_client_message({"type": "snapshot_state"})
|
|
assert isinstance(parsed, SnapshotState)
|
|
|
|
def test_unknown_type_rejected(self):
|
|
with pytest.raises(ValidationError):
|
|
parse_client_message({"type": "not_a_real_message"})
|
|
|
|
def test_missing_type_rejected(self):
|
|
with pytest.raises(ValidationError):
|
|
parse_client_message({"prompt": "x"})
|
|
|
|
def test_segment_prompt_source_requires_prompt(self):
|
|
with pytest.raises(ValidationError):
|
|
parse_client_message({"type": "segment_prompt_source"})
|
|
|
|
|
|
class TestServerMessageSerialization:
|
|
|
|
def test_queue_status(self):
|
|
msg = QueueStatus(position=3, queue_depth=5)
|
|
assert msg.model_dump() == {
|
|
"type": "queue_status",
|
|
"position": 3,
|
|
"queue_depth": 5,
|
|
}
|
|
|
|
def test_gpu_assigned(self):
|
|
msg = GpuAssigned(gpu_id=1, session_timeout=300)
|
|
assert msg.model_dump()["type"] == "gpu_assigned"
|
|
|
|
def test_ltx2_stream_start(self):
|
|
msg = Ltx2StreamStart(
|
|
preset="ltx2_two_stage",
|
|
width=1024,
|
|
height=1536,
|
|
fps=24,
|
|
num_frames=121,
|
|
)
|
|
dumped = msg.model_dump()
|
|
assert dumped["type"] == "ltx2_stream_start"
|
|
assert dumped["width"] == 1024
|
|
|
|
def test_ltx2_segment_start(self):
|
|
msg = Ltx2SegmentStart(
|
|
segment_idx=0,
|
|
prompt="a fox",
|
|
total_steps=8,
|
|
)
|
|
assert msg.model_dump()["segment_idx"] == 0
|
|
|
|
def test_step_complete(self):
|
|
msg = StepComplete(segment_idx=0, step=1, total_steps=8)
|
|
assert msg.model_dump()["stage"] == "denoise"
|
|
|
|
def test_media_init_has_mode(self):
|
|
msg = MediaInit(segment_idx=0, stream_id="abc")
|
|
dumped = msg.model_dump()
|
|
assert dumped["mode"] == "av_fmp4"
|
|
assert "avc1" in dumped["mime"]
|
|
|
|
def test_media_segment_complete(self):
|
|
msg = MediaSegmentComplete(
|
|
segment_idx=0,
|
|
stream_id="abc",
|
|
chunks=4,
|
|
)
|
|
dumped = msg.model_dump()
|
|
assert dumped["chunks"] == 4
|
|
|
|
def test_ltx2_segment_complete(self):
|
|
msg = Ltx2SegmentComplete(segment_idx=0, generation_time_ms=1234.5)
|
|
assert msg.model_dump()["generation_time_ms"] == 1234.5
|
|
|
|
def test_error_message_code_restricted(self):
|
|
with pytest.raises(ValidationError):
|
|
ErrorMessage(code="not_a_code", message="x")
|
|
|
|
def test_continuation_state_snapshot(self):
|
|
msg = ContinuationStateSnapshot(state={
|
|
"kind": "ltx2.v1",
|
|
"payload": {
|
|
"schema_version": 1
|
|
},
|
|
})
|
|
assert msg.model_dump()["state"]["kind"] == "ltx2.v1"
|