228 lines
7.8 KiB
Python
228 lines
7.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""End-to-end WebSocket smoke for the streaming server skeleton.
|
|
|
|
Uses a mock generator so these tests run CPU-only (no GPU, no model
|
|
weights). Skips the fMP4 assertions when ``ffmpeg`` is missing.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
pytest.importorskip("starlette")
|
|
from starlette.testclient import TestClient # noqa: E402
|
|
|
|
from fastvideo.api.schema import ( # noqa: E402
|
|
ContinuationState, GeneratorConfig, SamplingConfig, ServeConfig, StreamingConfig, GenerationRequest,
|
|
)
|
|
from fastvideo.entrypoints.streaming.server import build_app # noqa: E402
|
|
|
|
_FFMPEG_AVAILABLE = shutil.which("ffmpeg") is not None
|
|
|
|
|
|
@dataclass
|
|
class _MockGenerator:
|
|
width: int = 64
|
|
height: int = 64
|
|
fps: int = 12
|
|
num_frames: int = 12
|
|
return_state: bool = True
|
|
|
|
def generate(self, request: GenerationRequest) -> dict[str, Any]:
|
|
frames = [np.full((self.height, self.width, 3), i * 5, dtype=np.uint8) for i in range(self.num_frames)]
|
|
state = (ContinuationState(
|
|
kind="ltx2.v1",
|
|
payload={
|
|
"schema_version": 1,
|
|
"segment_index": 0,
|
|
"source_prompt": request.prompt,
|
|
},
|
|
) if self.return_state else None)
|
|
return {
|
|
"frames": frames,
|
|
"audio_sample_rate": 24000,
|
|
"state": state,
|
|
}
|
|
|
|
|
|
def _build_serve_config() -> ServeConfig:
|
|
return ServeConfig(
|
|
generator=GeneratorConfig(model_path="/models/fake"),
|
|
default_request=GenerationRequest(sampling=SamplingConfig(
|
|
num_frames=12,
|
|
height=64,
|
|
width=64,
|
|
fps=12,
|
|
num_inference_steps=1,
|
|
), ),
|
|
streaming=StreamingConfig(
|
|
session_timeout_seconds=60,
|
|
generation_segment_cap=2,
|
|
),
|
|
)
|
|
|
|
|
|
def _build_client() -> tuple[TestClient, _MockGenerator]:
|
|
generator = _MockGenerator()
|
|
app = build_app(_build_serve_config(), generator)
|
|
return TestClient(app), generator
|
|
|
|
|
|
class TestHealth:
|
|
|
|
def test_health_endpoint_reports_stream_mode(self):
|
|
client, _ = _build_client()
|
|
response = client.get("/health")
|
|
assert response.status_code == 200
|
|
body = response.json()
|
|
assert body["status"] == "ok"
|
|
assert body["stream_mode"] == "av_fmp4"
|
|
assert body["sessions"] == 0
|
|
|
|
|
|
class TestSessionHandshake:
|
|
|
|
def test_rejects_non_session_init_opening_frame(self):
|
|
client, _ = _build_client()
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
ws.send_json({"type": "segment_prompt_source", "prompt": "x"})
|
|
err = ws.receive_json()
|
|
assert err["type"] == "error"
|
|
assert err["code"] == "invalid_message"
|
|
|
|
def test_rejects_unknown_message_on_init(self):
|
|
client, _ = _build_client()
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
ws.send_json({"type": "not_a_message"})
|
|
err = ws.receive_json()
|
|
assert err["type"] == "error"
|
|
|
|
def test_emits_queue_and_gpu_assigned_on_valid_init(self):
|
|
client, _ = _build_client()
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
ws.send_json({
|
|
"type": "session_init_v2",
|
|
"preset": "ltx2_two_stage",
|
|
"curated_prompts": ["a fox"],
|
|
})
|
|
assert ws.receive_json()["type"] == "queue_status"
|
|
assert ws.receive_json()["type"] == "gpu_assigned"
|
|
assert ws.receive_json()["type"] == "ltx2_stream_start"
|
|
|
|
def test_init_hydrates_continuation_state(self):
|
|
client, _ = _build_client()
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
ws.send_json({
|
|
"type": "session_init_v2",
|
|
"preset": "ltx2_two_stage",
|
|
"continuation_state": {
|
|
"kind": "ltx2.v1",
|
|
"payload": {
|
|
"schema_version": 1,
|
|
"segment_index": 3
|
|
},
|
|
},
|
|
})
|
|
# Drain handshake frames
|
|
ws.receive_json() # queue_status
|
|
ws.receive_json() # gpu_assigned
|
|
ws.receive_json() # ltx2_stream_start
|
|
# Ask the server for the state back; it should echo what we sent.
|
|
ws.send_json({"type": "snapshot_state"})
|
|
snap = ws.receive_json()
|
|
assert snap["type"] == "continuation_state_snapshot"
|
|
assert snap["state"]["kind"] == "ltx2.v1"
|
|
assert snap["state"]["payload"]["segment_index"] == 3
|
|
|
|
|
|
@pytest.mark.skipif(not _FFMPEG_AVAILABLE, reason="ffmpeg not installed")
|
|
class TestSegmentFlow:
|
|
|
|
def test_segment_generates_media_init_plus_complete(self):
|
|
client, generator = _build_client()
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"})
|
|
for _ in range(3):
|
|
ws.receive_json() # queue_status + gpu_assigned + stream_start
|
|
|
|
ws.send_json({
|
|
"type": "segment_prompt_source",
|
|
"prompt": "a test segment",
|
|
"num_inference_steps": 1,
|
|
})
|
|
start = ws.receive_json()
|
|
assert start["type"] == "ltx2_segment_start"
|
|
assert start["segment_idx"] == 0
|
|
step = ws.receive_json()
|
|
assert step["type"] == "step_complete"
|
|
media_init = ws.receive_json()
|
|
assert media_init["type"] == "media_init"
|
|
# Then one or more binary frames until media_segment_complete.
|
|
saw_binary = False
|
|
while True:
|
|
msg = ws.receive()
|
|
if "bytes" in msg and msg["bytes"]:
|
|
saw_binary = True
|
|
continue
|
|
parsed = _as_json(msg)
|
|
if parsed is None:
|
|
continue
|
|
if parsed["type"] == "media_segment_complete":
|
|
break
|
|
assert saw_binary
|
|
final = ws.receive_json()
|
|
assert final["type"] == "ltx2_segment_complete"
|
|
assert final["segment_idx"] == 0
|
|
|
|
|
|
class TestContinuationStatePersistence:
|
|
|
|
def test_snapshot_after_segment_carries_generator_state(self):
|
|
if not _FFMPEG_AVAILABLE:
|
|
pytest.skip("ffmpeg not installed")
|
|
client, generator = _build_client()
|
|
with client.websocket_connect("/v1/stream") as ws:
|
|
ws.send_json({"type": "session_init_v2", "preset": "ltx2_two_stage"})
|
|
for _ in range(3):
|
|
ws.receive_json()
|
|
ws.send_json({
|
|
"type": "segment_prompt_source",
|
|
"prompt": "a cat",
|
|
"num_inference_steps": 1,
|
|
})
|
|
_drain_until(ws, "ltx2_segment_complete")
|
|
ws.send_json({"type": "snapshot_state"})
|
|
snap = ws.receive_json()
|
|
assert snap["type"] == "continuation_state_snapshot"
|
|
assert snap["state"]["kind"] == "ltx2.v1"
|
|
assert snap["state"]["payload"]["source_prompt"] == "a cat"
|
|
|
|
|
|
# ----------------------------------------------------------------------
|
|
# Helpers
|
|
# ----------------------------------------------------------------------
|
|
|
|
|
|
def _drain_until(ws, target_type: str) -> dict[str, Any]:
|
|
while True:
|
|
msg = ws.receive()
|
|
if "text" in msg and msg["text"]:
|
|
import json
|
|
|
|
parsed = json.loads(msg["text"])
|
|
if parsed.get("type") == target_type:
|
|
return parsed
|
|
# skip binary / other
|
|
|
|
|
|
def _as_json(msg: dict[str, Any]) -> dict[str, Any] | None:
|
|
if "text" not in msg or not msg["text"]:
|
|
return None
|
|
import json
|
|
|
|
return json.loads(msg["text"])
|