Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: shimei.jmh <shimei.jmh@taobao.com> Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
This commit is contained in:
14 files changed
+3597
No files matched your search
@@ -0,0 +1,31 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Shared test helpers for video streaming tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
|
||||
import pytest
|
||||
|
||||
np = pytest.importorskip("numpy", reason="numpy required for video stream tests")
|
||||
PIL = pytest.importorskip("PIL", reason="Pillow required for video stream tests")
|
||||
from PIL import Image # noqa: E402
|
||||
|
||||
|
||||
def make_jpeg(r: int = 128, g: int = 128, b: int = 128, size: int = 64) -> bytes:
|
||||
"""Create a solid-colour JPEG image."""
|
||||
img = Image.new("RGB", (size, size), (r, g, b))
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=95)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def make_gradient_jpeg(seed: int, size: int = 64) -> bytes:
|
||||
"""Create a random-gradient JPEG that varies based on *seed*."""
|
||||
rng = np.random.RandomState(seed)
|
||||
arr = rng.randint(0, 256, (size, size, 3), dtype=np.uint8)
|
||||
img = Image.fromarray(arr, "RGB")
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=95)
|
||||
return buf.getvalue()
|
||||
@@ -0,0 +1,650 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the serving-layer streaming video WebSocket handler."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.entrypoints.openai import serving_video_stream, video_stream_envs
|
||||
from vllm_omni.entrypoints.openai.serving_video_stream import (
|
||||
OmniStreamingVideoHandler,
|
||||
StreamingVideoSessionConfig,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes:
|
||||
img = Image.new("RGB", (64, 64), (r, g, b))
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=95)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _b64(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _text_result(text: str) -> OmniRequestOutput:
|
||||
class Output:
|
||||
pass
|
||||
|
||||
class RequestOutput:
|
||||
pass
|
||||
|
||||
output = Output()
|
||||
output.text = text
|
||||
request_output = RequestOutput()
|
||||
request_output.outputs = [output]
|
||||
return OmniRequestOutput(final_output_type="text", request_output=request_output)
|
||||
|
||||
|
||||
def _audio_result(audio_data: Any) -> OmniRequestOutput:
|
||||
class Output:
|
||||
pass
|
||||
|
||||
class RequestOutput:
|
||||
pass
|
||||
|
||||
output = Output()
|
||||
output.multimodal_output = {"audio": audio_data}
|
||||
request_output = RequestOutput()
|
||||
request_output.outputs = [output]
|
||||
return OmniRequestOutput(final_output_type="audio", request_output=request_output)
|
||||
|
||||
|
||||
class MockWebSocket:
|
||||
def __init__(self, messages: list[str] | None = None):
|
||||
self._messages = list(messages or [])
|
||||
self._idx = 0
|
||||
self.accepted = False
|
||||
self.sent: list[dict[str, Any]] = []
|
||||
|
||||
async def accept(self):
|
||||
self.accepted = True
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
if self._idx >= len(self._messages):
|
||||
await asyncio.sleep(999)
|
||||
msg = self._messages[self._idx]
|
||||
self._idx += 1
|
||||
return msg
|
||||
|
||||
async def send_json(self, data: dict[str, Any]):
|
||||
self.sent.append(data)
|
||||
|
||||
|
||||
class TimedWebSocket:
|
||||
def __init__(self):
|
||||
self._q: asyncio.Queue[str] = asyncio.Queue()
|
||||
self.accepted = False
|
||||
self.sent: list[dict[str, Any]] = []
|
||||
|
||||
async def accept(self):
|
||||
self.accepted = True
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
return await self._q.get()
|
||||
|
||||
async def send_json(self, data: dict[str, Any]):
|
||||
self.sent.append(data)
|
||||
|
||||
def put(self, msg: dict[str, Any]):
|
||||
self._q.put_nowait(json.dumps(msg))
|
||||
|
||||
def sent_types(self) -> list[str]:
|
||||
return [m.get("type", "") for m in self.sent]
|
||||
|
||||
|
||||
def test_api_server_registers_video_stream_route():
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
|
||||
assert any(getattr(route, "path", None) == "/v1/video/chat/stream" for route in router.routes)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_receive_config_accepts_client_legacy_aliases():
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": "test",
|
||||
"num_sample_frames": 7,
|
||||
"evs_enabled": False,
|
||||
"evs_threshold": 0.87,
|
||||
}
|
||||
)
|
||||
]
|
||||
)
|
||||
handler = OmniStreamingVideoHandler(chat_service=object())
|
||||
|
||||
config = await handler._receive_config(ws)
|
||||
|
||||
assert config is not None
|
||||
assert config.num_frames == 7
|
||||
assert config.enable_frame_filter is False
|
||||
assert config.frame_filter_threshold == 0.87
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audio_in_video_sets_mm_processor_kwargs():
|
||||
captured_requests = []
|
||||
|
||||
class EmptyEngine:
|
||||
def generate(self, **_kwargs):
|
||||
async def _gen():
|
||||
if False:
|
||||
yield None
|
||||
|
||||
return _gen()
|
||||
|
||||
class CapturingHandler(OmniStreamingVideoHandler):
|
||||
async def _preprocess_to_engine_prompt(self, request):
|
||||
captured_requests.append(request)
|
||||
return {"prompt": "x"}
|
||||
|
||||
ws = MockWebSocket()
|
||||
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine())
|
||||
config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=True)
|
||||
|
||||
await handler._process_query_engine(
|
||||
ws,
|
||||
config,
|
||||
[_b64(_make_jpeg())],
|
||||
bytearray(b"\x00\x00"),
|
||||
[],
|
||||
"what is happening?",
|
||||
"req-1",
|
||||
asyncio.Event(),
|
||||
{},
|
||||
)
|
||||
|
||||
assert captured_requests
|
||||
assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audio_in_video_disabled_omits_mm_processor_kwargs():
|
||||
captured_requests = []
|
||||
|
||||
class EmptyEngine:
|
||||
def generate(self, **_kwargs):
|
||||
async def _gen():
|
||||
if False:
|
||||
yield None
|
||||
|
||||
return _gen()
|
||||
|
||||
class CapturingHandler(OmniStreamingVideoHandler):
|
||||
async def _preprocess_to_engine_prompt(self, request):
|
||||
captured_requests.append(request)
|
||||
return {"prompt": "x"}
|
||||
|
||||
ws = MockWebSocket()
|
||||
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine())
|
||||
config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=False)
|
||||
|
||||
await handler._process_query_engine(
|
||||
ws,
|
||||
config,
|
||||
[_b64(_make_jpeg())],
|
||||
bytearray(b"\x00\x00"),
|
||||
[],
|
||||
"what is happening?",
|
||||
"req-1",
|
||||
asyncio.Event(),
|
||||
{},
|
||||
)
|
||||
|
||||
assert captured_requests
|
||||
assert captured_requests[0].mm_processor_kwargs is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_inline_audio_data_sets_mm_processor_kwargs():
|
||||
captured_requests = []
|
||||
|
||||
class EmptyEngine:
|
||||
def generate(self, **_kwargs):
|
||||
async def _gen():
|
||||
if False:
|
||||
yield None
|
||||
|
||||
return _gen()
|
||||
|
||||
class CapturingHandler(OmniStreamingVideoHandler):
|
||||
async def _preprocess_to_engine_prompt(self, request):
|
||||
captured_requests.append(request)
|
||||
return {"prompt": "x"}
|
||||
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
json.dumps({"type": "session.config", "model": "test"}),
|
||||
json.dumps({"type": "video.frame", "data": _b64(_make_jpeg())}),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "video.query",
|
||||
"text": "describe",
|
||||
"audio_data": _b64(b"\x00\x00"),
|
||||
}
|
||||
),
|
||||
json.dumps({"type": "video.done"}),
|
||||
]
|
||||
)
|
||||
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=2.0)
|
||||
|
||||
await handler.handle_session(ws)
|
||||
|
||||
assert captured_requests
|
||||
assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True}
|
||||
assert "session.done" in [m.get("type") for m in ws.sent]
|
||||
|
||||
|
||||
def test_audio_delta_mode_is_read_by_serving_code_at_runtime(monkeypatch):
|
||||
handler = OmniStreamingVideoHandler(chat_service=object())
|
||||
result = _audio_result([object()])
|
||||
|
||||
monkeypatch.setattr(
|
||||
OmniStreamingVideoHandler,
|
||||
"_delta_fast",
|
||||
classmethod(lambda cls, audio_data, chunks_drained: ("fast-path", chunks_drained)),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
OmniStreamingVideoHandler,
|
||||
"_delta_slow",
|
||||
classmethod(lambda cls, audio_data, chunks_drained: ("slow-path", chunks_drained)),
|
||||
)
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "fast")
|
||||
assert handler._extract_audio_delta_b64(result, 0)[0] == "fast-path"
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "slow")
|
||||
assert handler._extract_audio_delta_b64(result, 0)[0] == "slow-path"
|
||||
|
||||
|
||||
def test_video_stream_envs_strip_and_warn_once_per_invalid_value(monkeypatch):
|
||||
warnings = []
|
||||
|
||||
video_stream_envs._warned_invalid_envs.clear()
|
||||
try:
|
||||
monkeypatch.setattr(
|
||||
video_stream_envs.logger,
|
||||
"warning",
|
||||
lambda message, *args, **_kwargs: warnings.append((message, args)),
|
||||
)
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", " off ")
|
||||
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "off"
|
||||
assert not warnings
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "bad")
|
||||
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on"
|
||||
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on"
|
||||
assert len(warnings) == 1
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "still_bad")
|
||||
assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on"
|
||||
assert len(warnings) == 2
|
||||
finally:
|
||||
video_stream_envs._warned_invalid_envs.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_chunk_mode_is_read_by_engine_path_at_runtime(monkeypatch):
|
||||
class TextEngine:
|
||||
def generate(self, **_kwargs):
|
||||
async def _gen():
|
||||
yield _text_result("hello")
|
||||
|
||||
return _gen()
|
||||
|
||||
class CapturingHandler(OmniStreamingVideoHandler):
|
||||
async def _preprocess_to_engine_prompt(self, request):
|
||||
return {"prompt": "x"}
|
||||
|
||||
handler = CapturingHandler(chat_service=object(), engine_client=TextEngine())
|
||||
config = StreamingVideoSessionConfig(model="test", modalities=["text"])
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "on")
|
||||
ws_on = MockWebSocket()
|
||||
await handler._process_query_engine(
|
||||
ws_on,
|
||||
config,
|
||||
[_b64(_make_jpeg())],
|
||||
bytearray(),
|
||||
[],
|
||||
"describe",
|
||||
"req-on",
|
||||
asyncio.Event(),
|
||||
{},
|
||||
)
|
||||
assert {"type": "response.text.delta", "delta": "hello"} in ws_on.sent
|
||||
|
||||
monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "off")
|
||||
ws_off = MockWebSocket()
|
||||
await handler._process_query_engine(
|
||||
ws_off,
|
||||
config,
|
||||
[_b64(_make_jpeg())],
|
||||
bytearray(),
|
||||
[],
|
||||
"describe",
|
||||
"req-off",
|
||||
asyncio.Event(),
|
||||
{},
|
||||
)
|
||||
assert {"type": "response.text.done", "text": "hello"} in ws_off.sent
|
||||
assert not any(m.get("type") == "response.text.delta" for m in ws_off.sent)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_without_engine_client_sends_error():
|
||||
ws = MockWebSocket()
|
||||
handler = OmniStreamingVideoHandler(chat_service=object(), engine_client=None)
|
||||
|
||||
await handler._process_query(
|
||||
ws,
|
||||
StreamingVideoSessionConfig(model="test"),
|
||||
[],
|
||||
bytearray(),
|
||||
[],
|
||||
"describe",
|
||||
"req-1",
|
||||
asyncio.Event(),
|
||||
{},
|
||||
)
|
||||
|
||||
assert {"type": "error", "message": "Streaming video requires an engine client"} in ws.sent
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_new_query_cancels_in_flight_query():
|
||||
query_started = asyncio.Event()
|
||||
query_cancelled = asyncio.Event()
|
||||
calls = 0
|
||||
|
||||
class BlockingHandler(OmniStreamingVideoHandler):
|
||||
async def _process_query(self, *args, **kwargs):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls > 1:
|
||||
return
|
||||
query_started.set()
|
||||
try:
|
||||
await asyncio.sleep(999)
|
||||
except asyncio.CancelledError:
|
||||
query_cancelled.set()
|
||||
raise
|
||||
|
||||
ws = TimedWebSocket()
|
||||
handler = BlockingHandler(chat_service=object(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
||||
|
||||
ws.put({"type": "video.query", "text": "interrupt"})
|
||||
await asyncio.wait_for(query_cancelled.wait(), timeout=2.0)
|
||||
ws.put({"type": "video.done"})
|
||||
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_video_done_waits_for_in_flight_query():
|
||||
query_started = asyncio.Event()
|
||||
allow_finish = asyncio.Event()
|
||||
query_finished = asyncio.Event()
|
||||
|
||||
class BlockingHandler(OmniStreamingVideoHandler):
|
||||
async def _process_query(self, *args, **kwargs):
|
||||
query_started.set()
|
||||
await allow_finish.wait()
|
||||
query_finished.set()
|
||||
|
||||
ws = TimedWebSocket()
|
||||
handler = BlockingHandler(chat_service=object(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
||||
|
||||
ws.put({"type": "video.done"})
|
||||
await asyncio.sleep(0.05)
|
||||
assert not task.done()
|
||||
assert not query_finished.is_set()
|
||||
|
||||
allow_finish.set()
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
|
||||
assert query_finished.is_set()
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_prewarm_does_not_block_following_query(monkeypatch):
|
||||
decode_started = threading.Event()
|
||||
release_decode = threading.Event()
|
||||
query_started = asyncio.Event()
|
||||
|
||||
def blocked_decode(raw_bytes: bytes):
|
||||
decode_started.set()
|
||||
release_decode.wait(timeout=2.0)
|
||||
return Image.open(io.BytesIO(raw_bytes)).convert("RGB")
|
||||
|
||||
class BlockingHandler(OmniStreamingVideoHandler):
|
||||
async def _process_query(self, *args, **kwargs):
|
||||
query_started.set()
|
||||
|
||||
monkeypatch.setattr(serving_video_stream, "_decode_frame_bytes", blocked_decode)
|
||||
|
||||
ws = TimedWebSocket()
|
||||
handler = BlockingHandler(chat_service=object(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
||||
|
||||
for _ in range(100):
|
||||
if decode_started.is_set():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
assert decode_started.is_set()
|
||||
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
||||
|
||||
release_decode.set()
|
||||
ws.put({"type": "video.done"})
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_cannot_send_internal_frame_decode_failed_message():
|
||||
captured_frames: list[list[str]] = []
|
||||
frame = _b64(_make_jpeg())
|
||||
|
||||
class CapturingHandler(OmniStreamingVideoHandler):
|
||||
async def _process_query(
|
||||
self,
|
||||
websocket,
|
||||
config,
|
||||
frame_buffer,
|
||||
audio_buffer,
|
||||
message_history,
|
||||
query_text,
|
||||
request_id,
|
||||
interrupt_event,
|
||||
prewarmed_frames,
|
||||
):
|
||||
captured_frames.append(list(frame_buffer))
|
||||
|
||||
ws = TimedWebSocket()
|
||||
handler = CapturingHandler(chat_service=object(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": frame})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "_internal.frame_decode_failed", "b64": frame})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.done"})
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
|
||||
assert {"type": "error", "message": "Unknown type: _internal.frame_decode_failed"} in ws.sent
|
||||
assert captured_frames == [[frame]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_frame_prewarm_removes_frame_before_query():
|
||||
ws = TimedWebSocket()
|
||||
handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test", "enable_frame_filter": False})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": _b64(b"not-a-jpeg")})
|
||||
|
||||
for _ in range(100):
|
||||
if any(m.get("message") == "Frame decode failed" for m in ws.sent):
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert {"type": "error", "message": "Frame decode failed"} in ws.sent
|
||||
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.done"})
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
|
||||
assert {"type": "error", "message": "No frames buffered"} in ws.sent
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_filter_error_sends_invalid_image(monkeypatch):
|
||||
def fail_should_retain(self, frame_jpeg):
|
||||
raise ValueError("decode failed")
|
||||
|
||||
monkeypatch.setattr(serving_video_stream.FrameSimilarityFilter, "should_retain", fail_should_retain)
|
||||
|
||||
ws = TimedWebSocket()
|
||||
handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.done"})
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
|
||||
assert {"type": "error", "message": "Invalid image data"} in ws.sent
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch):
|
||||
captured_audio_lengths: list[int] = []
|
||||
|
||||
class EmptyEngine:
|
||||
def generate(self, **_kwargs):
|
||||
async def _gen():
|
||||
if False:
|
||||
yield None
|
||||
|
||||
return _gen()
|
||||
|
||||
class CapturingHandler(OmniStreamingVideoHandler):
|
||||
async def _process_query_engine(
|
||||
self,
|
||||
websocket,
|
||||
config,
|
||||
frame_buffer,
|
||||
audio_buffer,
|
||||
message_history,
|
||||
query_text,
|
||||
request_id,
|
||||
interrupt_event,
|
||||
prewarmed_frames,
|
||||
):
|
||||
captured_audio_lengths.append(len(audio_buffer))
|
||||
|
||||
monkeypatch.setattr(serving_video_stream, "_MAX_AUDIO_BUFFER_BYTES", 4)
|
||||
|
||||
ws = TimedWebSocket()
|
||||
handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=5.0)
|
||||
task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
ws.put({"type": "session.config", "model": "test"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "audio.chunk", "data": _b64(b"1234")})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "audio.chunk", "data": _b64(b"5")})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.frame", "data": _b64(_make_jpeg())})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.query", "text": "describe"})
|
||||
await asyncio.sleep(0)
|
||||
ws.put({"type": "video.done"})
|
||||
await asyncio.wait_for(task, timeout=2.0)
|
||||
|
||||
assert {"type": "error", "message": "Audio buffer overflow"} in ws.sent
|
||||
assert captured_audio_lengths == [0]
|
||||
|
||||
|
||||
def test_build_messages_keeps_recent_history_text_only():
|
||||
handler = OmniStreamingVideoHandler(chat_service=object())
|
||||
old_frame = _b64(_make_jpeg(1, 2, 3))
|
||||
current_frame = _b64(_make_jpeg(4, 5, 6))
|
||||
history = [
|
||||
{"role": "user", "content": [{"type": "text", "text": "old question"}]},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{old_frame}"}},
|
||||
{"type": "input_audio", "input_audio": {"data": "ignored", "format": "wav"}},
|
||||
{"type": "text", "text": "recent question"},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "recent answer"},
|
||||
]
|
||||
|
||||
messages, user_message = handler._build_messages(
|
||||
StreamingVideoSessionConfig(model="test", num_frames=1),
|
||||
[current_frame],
|
||||
bytearray(),
|
||||
history,
|
||||
"current question",
|
||||
{},
|
||||
)
|
||||
|
||||
assert messages[0] == {"role": "user", "content": "recent question"}
|
||||
assert messages[1] == {"role": "assistant", "content": "recent answer"}
|
||||
assert messages[2] == user_message
|
||||
assert user_message["content"][-1] == {"type": "text", "text": "current question"}
|
||||
@@ -0,0 +1,89 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for FrameSimilarityFilter (Phase 2 EVS)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai_api.conftest_video import (
|
||||
make_gradient_jpeg,
|
||||
make_jpeg,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter
|
||||
|
||||
|
||||
class TestFrameSimilarityFilter:
|
||||
def test_first_frame_always_retained(self):
|
||||
f = FrameSimilarityFilter(threshold=0.99)
|
||||
assert f.should_retain(make_jpeg()) is True
|
||||
|
||||
def test_identical_frames_dropped(self):
|
||||
f = FrameSimilarityFilter(threshold=0.90)
|
||||
frame = make_jpeg(100, 100, 100)
|
||||
assert f.should_retain(frame) is True
|
||||
assert f.should_retain(frame) is False
|
||||
|
||||
def test_very_different_frames_retained(self):
|
||||
f = FrameSimilarityFilter(threshold=0.95)
|
||||
assert f.should_retain(make_jpeg(255, 255, 255)) is True
|
||||
assert f.should_retain(make_jpeg(0, 0, 0)) is True
|
||||
|
||||
def test_low_threshold_keeps_slightly_different(self):
|
||||
f = FrameSimilarityFilter(threshold=0.50)
|
||||
assert f.should_retain(make_jpeg(100, 100, 100)) is True
|
||||
assert f.should_retain(make_jpeg(110, 110, 110)) is True
|
||||
|
||||
def test_random_frames_all_retained(self):
|
||||
f = FrameSimilarityFilter(threshold=0.95)
|
||||
for i in range(5):
|
||||
assert f.should_retain(make_gradient_jpeg(seed=i)) is True
|
||||
|
||||
def test_reset_clears_state(self):
|
||||
f = FrameSimilarityFilter(threshold=0.90)
|
||||
frame = make_jpeg(128, 128, 128)
|
||||
assert f.should_retain(frame) is True
|
||||
assert f.should_retain(frame) is False
|
||||
f.reset()
|
||||
assert f.should_retain(frame) is True
|
||||
|
||||
def test_stats_counting(self):
|
||||
f = FrameSimilarityFilter(threshold=0.90)
|
||||
frame = make_jpeg(50, 50, 50)
|
||||
f.should_retain(frame) # retained
|
||||
f.should_retain(frame) # dropped
|
||||
f.should_retain(frame) # dropped
|
||||
stats = f.stats
|
||||
assert stats["retained_count"] == 1
|
||||
assert stats["dropped_count"] == 2
|
||||
assert stats["total_count"] == 3
|
||||
assert abs(stats["drop_rate"] - 2.0 / 3.0) < 1e-6
|
||||
|
||||
def test_stats_empty(self):
|
||||
stats = FrameSimilarityFilter().stats
|
||||
assert stats["total_count"] == 0
|
||||
assert stats["drop_rate"] == 0.0
|
||||
|
||||
def test_stats_reset(self):
|
||||
f = FrameSimilarityFilter(threshold=0.90)
|
||||
f.should_retain(make_jpeg())
|
||||
f.reset()
|
||||
assert f.stats["total_count"] == 0
|
||||
|
||||
def test_invalid_threshold(self):
|
||||
with pytest.raises(ValueError, match="threshold"):
|
||||
FrameSimilarityFilter(threshold=1.5)
|
||||
with pytest.raises(ValueError, match="threshold"):
|
||||
FrameSimilarityFilter(threshold=-0.1)
|
||||
|
||||
def test_invalid_thumbnail_size(self):
|
||||
with pytest.raises(ValueError, match="thumbnail_size"):
|
||||
FrameSimilarityFilter(thumbnail_size=0)
|
||||
|
||||
def test_different_image_sizes_same_colour(self):
|
||||
"""Filter should handle frames of varying resolutions."""
|
||||
f = FrameSimilarityFilter(threshold=0.90)
|
||||
small = make_jpeg(100, 100, 100, size=32)
|
||||
large = make_jpeg(100, 100, 100, size=256)
|
||||
assert f.should_retain(small) is True
|
||||
assert f.should_retain(large) is False
|
||||
@@ -0,0 +1,593 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Runtime behavior tests for VideoStreamHandler.
|
||||
|
||||
Every test creates a mock WebSocket, drives handle_session through a
|
||||
specific code path, and asserts on the JSON messages actually sent back.
|
||||
No inspect.getsource() tricks — these test real async execution.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.entrypoints.openai.video_stream_session import (
|
||||
VideoStreamHandler,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes:
|
||||
img = Image.new("RGB", (64, 64), (r, g, b))
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="JPEG", quality=95)
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def _b64(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode()
|
||||
|
||||
|
||||
def _config_msg(**overrides) -> str:
|
||||
msg = {"type": "session.config", "model": "test", "evs_enabled": False}
|
||||
msg.update(overrides)
|
||||
return json.dumps(msg)
|
||||
|
||||
|
||||
def _frame_msg(jpeg: bytes | None = None) -> str:
|
||||
if jpeg is None:
|
||||
jpeg = _make_jpeg()
|
||||
return json.dumps({"type": "video.frame", "data": _b64(jpeg)})
|
||||
|
||||
|
||||
def _query_msg(text: str = "What do you see?") -> str:
|
||||
return json.dumps({"type": "video.query", "text": text})
|
||||
|
||||
|
||||
def _done_msg() -> str:
|
||||
return json.dumps({"type": "video.done"})
|
||||
|
||||
|
||||
def _audio_msg(pcm: bytes = b"\x00" * 320) -> str:
|
||||
return json.dumps({"type": "audio.chunk", "data": _b64(pcm)})
|
||||
|
||||
|
||||
def _make_sse_line(content: str) -> str:
|
||||
payload = {"choices": [{"delta": {"content": content}}]}
|
||||
return f"data: {json.dumps(payload)}\n\n"
|
||||
|
||||
|
||||
class MockWebSocket:
|
||||
"""Async-compatible mock WebSocket that feeds messages from a list."""
|
||||
|
||||
def __init__(self, messages: list[str]):
|
||||
self._messages = list(messages)
|
||||
self._idx = 0
|
||||
self._sent: list[dict | bytes] = []
|
||||
self._accepted = False
|
||||
|
||||
async def accept(self):
|
||||
self._accepted = True
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
if self._idx >= len(self._messages):
|
||||
# Simulate connection close by hanging forever (will be timed out)
|
||||
await asyncio.sleep(999)
|
||||
msg = self._messages[self._idx]
|
||||
self._idx += 1
|
||||
return msg
|
||||
|
||||
async def send_json(self, data: dict):
|
||||
self._sent.append(data)
|
||||
|
||||
async def send_bytes(self, data: bytes):
|
||||
self._sent.append(data)
|
||||
|
||||
@property
|
||||
def sent_messages(self) -> list[dict]:
|
||||
return [m for m in self._sent if isinstance(m, dict)]
|
||||
|
||||
def sent_types(self) -> list[str]:
|
||||
return [m.get("type", "") for m in self.sent_messages]
|
||||
|
||||
|
||||
class MockChatHandler:
|
||||
"""Mock OmniOpenAIServingChat for testing _handle_query paths."""
|
||||
|
||||
def __init__(self, response=None):
|
||||
self._response = response
|
||||
|
||||
async def create_chat_completion(self, request, raw_request=None):
|
||||
if self._response is not None:
|
||||
return self._response
|
||||
|
||||
# Default: return a simple streaming generator
|
||||
async def _gen():
|
||||
yield _make_sse_line("Hello")
|
||||
yield _make_sse_line("World")
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return _gen()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _receive_config path tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestReceiveConfig:
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_timeout(self):
|
||||
"""No message within config_timeout → error sent, session ends."""
|
||||
ws = MockWebSocket([]) # no messages — will hang
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=0.05)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
assert ws._accepted
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert len(errors) == 1
|
||||
assert "Timeout" in errors[0]["message"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_invalid_json(self):
|
||||
"""Non-JSON config message → error, session ends."""
|
||||
ws = MockWebSocket(["not json"])
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Invalid JSON" in e["message"] for e in errors)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_wrong_type(self):
|
||||
"""Message with wrong type → error, session ends."""
|
||||
ws = MockWebSocket([json.dumps({"type": "wrong"})])
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Expected session.config" in e["message"] for e in errors)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_config_invalid_field_type(self):
|
||||
"""Config with wrong field type → error, session ends."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.config",
|
||||
"model": "test",
|
||||
"max_frames": "potato",
|
||||
}
|
||||
)
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("max_frames" in e["message"] for e in errors)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Normal session flow
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestNormalFlow:
|
||||
@pytest.mark.asyncio
|
||||
async def test_frame_query_done(self):
|
||||
"""Happy path: config → frame → query → done → session.done."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
_frame_msg(),
|
||||
_query_msg("Describe."),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
types = ws.sent_types()
|
||||
assert "response.start" in types
|
||||
assert "response.text.delta" in types
|
||||
assert "response.text.done" in types
|
||||
assert "session.done" in types
|
||||
|
||||
# Check text content was streamed
|
||||
deltas = [m["delta"] for m in ws.sent_messages if m.get("type") == "response.text.delta"]
|
||||
assert "Hello" in deltas
|
||||
assert "World" in deltas
|
||||
|
||||
# response.text.done has full text
|
||||
done_msg = next(m for m in ws.sent_messages if m.get("type") == "response.text.done")
|
||||
assert done_msg["text"] == "HelloWorld"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_frames_before_query(self):
|
||||
"""Multiple frames accumulate, query uses sampled frames."""
|
||||
frames = [_frame_msg(_make_jpeg(r=i * 50)) for i in range(5)]
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(num_sample_frames=3),
|
||||
*frames,
|
||||
_query_msg(),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_evs_stats_in_done(self):
|
||||
"""EVS stats sent before session.done when evs_enabled=True."""
|
||||
frame = _frame_msg(_make_jpeg(100, 100, 100))
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(evs_enabled=True, evs_threshold=0.90),
|
||||
frame,
|
||||
frame, # second should be dropped
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
evs_msgs = [m for m in ws.sent_messages if m.get("type") == "response.evs_stats"]
|
||||
assert len(evs_msgs) == 1
|
||||
assert evs_msgs[0]["retained_count"] == 1
|
||||
assert evs_msgs[0]["dropped_count"] == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _handle_query error paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestHandleQueryErrors:
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_with_no_frames(self):
|
||||
"""Query before any frames → error, session continues."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
_query_msg("Hello?"),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("No frames" in e["message"] for e in errors)
|
||||
# Session still ends cleanly
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_returns_error_response(self):
|
||||
"""create_chat_completion returns ErrorResponse → error forwarded."""
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
|
||||
|
||||
err = ErrorResponse(
|
||||
message="model overloaded",
|
||||
type="server_error",
|
||||
code=503,
|
||||
)
|
||||
|
||||
chat = AsyncMock()
|
||||
chat.create_chat_completion = AsyncMock(return_value=err)
|
||||
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
_frame_msg(),
|
||||
_query_msg(),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=chat, idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
# Should see response.start, then error, then response.text.done
|
||||
assert "response.start" in ws.sent_types()
|
||||
assert "response.text.done" in ws.sent_types()
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("model overloaded" in e["message"] for e in errors)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_query_generator_raises(self):
|
||||
"""Exception during streaming → error sent, session continues."""
|
||||
|
||||
async def _exploding_gen():
|
||||
yield _make_sse_line("partial")
|
||||
raise RuntimeError("CUDA OOM")
|
||||
|
||||
chat = MockChatHandler()
|
||||
chat._response = None # override below
|
||||
|
||||
class BoomChat:
|
||||
async def create_chat_completion(self, req, raw_request=None):
|
||||
return _exploding_gen()
|
||||
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
_frame_msg(),
|
||||
_query_msg(),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=BoomChat(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Query processing failed" in e["message"] for e in errors)
|
||||
# Crucially: no raw exception message leaked
|
||||
assert not any("CUDA OOM" in e["message"] for e in errors)
|
||||
# Protocol: response.text.done always sent after response.start
|
||||
assert "response.text.done" in ws.sent_types()
|
||||
# Session still ends cleanly
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_query_text_rejected(self):
|
||||
"""video.query with empty text → error, session continues."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
_frame_msg(),
|
||||
json.dumps({"type": "video.query", "text": ""}),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("non-empty" in e["message"] for e in errors)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Timeout paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestTimeouts:
|
||||
@pytest.mark.asyncio
|
||||
async def test_idle_timeout(self):
|
||||
"""No message after config within idle_timeout → error, session ends."""
|
||||
ws = MockWebSocket([_config_msg()]) # config then nothing
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=0.05)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Idle timeout" in e["message"] for e in errors)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reader error paths
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestReaderErrors:
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_json_mid_session(self):
|
||||
"""Invalid JSON mid-session → error sent, session continues."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
"{{bad json",
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Invalid JSON" in e["message"] for e in errors)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_message(self):
|
||||
"""Non-object JSON → error sent, continues."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
json.dumps([1, 2, 3]),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("JSON objects" in e["message"] for e in errors)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_message_type(self):
|
||||
"""Unknown type → error, session continues."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
json.dumps({"type": "teleport.now"}),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Unknown message type" in e["message"] for e in errors)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_base64_frame(self):
|
||||
"""Invalid base64 in video.frame → error, continues."""
|
||||
ws = MockWebSocket(
|
||||
[
|
||||
_config_msg(),
|
||||
json.dumps({"type": "video.frame", "data": "!!!not-base64!!!"}),
|
||||
_done_msg(),
|
||||
]
|
||||
)
|
||||
handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0)
|
||||
await handler.handle_session(ws)
|
||||
|
||||
errors = [m for m in ws.sent_messages if m["type"] == "error"]
|
||||
assert any("Invalid base64" in e["message"] for e in errors)
|
||||
assert "session.done" in ws.sent_types()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Concurrency: frames arrive during query processing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestConcurrency:
|
||||
@pytest.mark.asyncio
|
||||
async def test_frames_during_query(self):
|
||||
"""Frames sent while a query is processing must be buffered."""
|
||||
query_started = asyncio.Event()
|
||||
query_can_finish = asyncio.Event()
|
||||
captured_frame_count: list[int] = []
|
||||
|
||||
class SlowChat:
|
||||
async def create_chat_completion(self, req, raw_request=None):
|
||||
query_started.set()
|
||||
await query_can_finish.wait()
|
||||
|
||||
async def _gen():
|
||||
yield _make_sse_line("answer")
|
||||
yield "data: [DONE]\n\n"
|
||||
|
||||
return _gen()
|
||||
|
||||
class InspectingSlowChat(SlowChat):
|
||||
"""On the second query, record how many frames are in the request."""
|
||||
|
||||
_call_count = 0
|
||||
|
||||
async def create_chat_completion(self, req, raw_request=None):
|
||||
self._call_count += 1
|
||||
if self._call_count == 2:
|
||||
# Count image_url blocks in the request to verify buffering
|
||||
for msg in req.messages:
|
||||
if isinstance(msg.get("content", msg), list):
|
||||
content = msg.get("content", msg)
|
||||
else:
|
||||
content = getattr(msg, "content", [])
|
||||
if isinstance(content, list):
|
||||
n = sum(
|
||||
1
|
||||
for p in content
|
||||
if (isinstance(p, dict) and p.get("type") == "image_url")
|
||||
or (hasattr(p, "type") and p.type == "image_url")
|
||||
)
|
||||
captured_frame_count.append(n)
|
||||
return await super().create_chat_completion(req, raw_request)
|
||||
|
||||
# Custom WebSocket that can inject frames after query starts
|
||||
class TimedWebSocket(MockWebSocket):
|
||||
def __init__(self):
|
||||
self._q: asyncio.Queue[str] = asyncio.Queue()
|
||||
self._sent: list[dict | bytes] = []
|
||||
self._accepted = False
|
||||
|
||||
async def accept(self):
|
||||
self._accepted = True
|
||||
|
||||
async def receive_text(self) -> str:
|
||||
return await self._q.get()
|
||||
|
||||
def put(self, msg: str):
|
||||
self._q.put_nowait(msg)
|
||||
|
||||
ws = TimedWebSocket()
|
||||
chat = InspectingSlowChat()
|
||||
handler = VideoStreamHandler(chat_handler=chat, idle_timeout=5.0)
|
||||
|
||||
session_task = asyncio.create_task(handler.handle_session(ws))
|
||||
|
||||
# 1. Config
|
||||
ws.put(_config_msg())
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
# 2. Initial frame + query
|
||||
ws.put(_frame_msg(_make_jpeg(10, 10, 10)))
|
||||
await asyncio.sleep(0.01)
|
||||
ws.put(_query_msg("first"))
|
||||
|
||||
# 3. Wait for query to start processing
|
||||
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
||||
|
||||
# 4. Send more frames WHILE query is running
|
||||
ws.put(_frame_msg(_make_jpeg(200, 200, 200)))
|
||||
ws.put(_frame_msg(_make_jpeg(50, 50, 50)))
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
# 5. Let query finish
|
||||
query_can_finish.set()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# 6. Second query — should see all 3 frames (1 initial + 2 during)
|
||||
query_started.clear()
|
||||
query_can_finish.clear()
|
||||
ws.put(_query_msg("second"))
|
||||
await asyncio.wait_for(query_started.wait(), timeout=2.0)
|
||||
query_can_finish.set()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
# 7. Done
|
||||
ws.put(_done_msg())
|
||||
|
||||
await asyncio.wait_for(session_task, timeout=5.0)
|
||||
|
||||
assert "session.done" in ws.sent_types()
|
||||
# Verify: the 2 frames sent during query 1 were buffered and
|
||||
# included in query 2 (total 3 frames).
|
||||
assert captured_frame_count == [3]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE parsing (kept from original, these are valid unit tests)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParseSSEDeltas:
|
||||
def test_single_delta(self):
|
||||
assert VideoStreamHandler._parse_sse_deltas(_make_sse_line("hello")) == ["hello"]
|
||||
|
||||
def test_multiple_deltas(self):
|
||||
chunk = _make_sse_line("a") + _make_sse_line("b")
|
||||
assert VideoStreamHandler._parse_sse_deltas(chunk) == ["a", "b"]
|
||||
|
||||
def test_done_skipped(self):
|
||||
chunk = _make_sse_line("x") + "data: [DONE]\n\n"
|
||||
assert VideoStreamHandler._parse_sse_deltas(chunk) == ["x"]
|
||||
|
||||
def test_empty_content_skipped(self):
|
||||
p = json.dumps({"choices": [{"delta": {"content": ""}}]})
|
||||
assert VideoStreamHandler._parse_sse_deltas(f"data: {p}\n\n") == []
|
||||
|
||||
def test_malformed_json_skipped(self):
|
||||
chunk = "data: {bad}\n\n" + _make_sse_line("ok")
|
||||
assert VideoStreamHandler._parse_sse_deltas(chunk) == ["ok"]
|
||||
|
||||
def test_empty_string(self):
|
||||
assert VideoStreamHandler._parse_sse_deltas("") == []
|
||||
|
||||
def test_unicode(self):
|
||||
assert VideoStreamHandler._parse_sse_deltas(_make_sse_line("你好🌍")) == ["你好🌍"]
|
||||
@@ -0,0 +1,289 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for VideoStreamSession (Phase 2 + 3)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.entrypoints.openai_api.conftest_video import (
|
||||
make_gradient_jpeg,
|
||||
make_jpeg,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.video_stream_session import (
|
||||
VideoStreamConfig,
|
||||
VideoStreamSession,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# VideoStreamConfig
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestVideoStreamConfig:
|
||||
def test_defaults(self):
|
||||
cfg = VideoStreamConfig()
|
||||
assert cfg.max_frames == 64
|
||||
assert cfg.num_sample_frames == 16
|
||||
assert cfg.evs_enabled is True
|
||||
assert cfg.evs_threshold == 0.95
|
||||
|
||||
def test_from_dict(self):
|
||||
cfg = VideoStreamConfig.from_dict(
|
||||
{
|
||||
"model": "test-model",
|
||||
"max_frames": 32,
|
||||
"evs_threshold": 0.90,
|
||||
"unknown_field": "ignored",
|
||||
}
|
||||
)
|
||||
assert cfg.model == "test-model"
|
||||
assert cfg.max_frames == 32
|
||||
assert cfg.evs_threshold == 0.90
|
||||
|
||||
def test_from_dict_empty(self):
|
||||
cfg = VideoStreamConfig.from_dict({})
|
||||
assert cfg.model == ""
|
||||
assert cfg.max_frames == 64
|
||||
|
||||
def test_from_dict_invalid_type(self):
|
||||
with pytest.raises(TypeError, match="max_frames.*expected int.*got str"):
|
||||
VideoStreamConfig.from_dict({"max_frames": "potato"})
|
||||
|
||||
def test_from_dict_invalid_bool(self):
|
||||
with pytest.raises(TypeError, match="evs_enabled.*expected bool"):
|
||||
VideoStreamConfig.from_dict({"evs_enabled": "yes"})
|
||||
|
||||
def test_from_dict_invalid_modalities(self):
|
||||
with pytest.raises(TypeError, match="modalities.*expected list"):
|
||||
VideoStreamConfig.from_dict({"modalities": 42})
|
||||
|
||||
def test_from_dict_evs_threshold_int_accepted(self):
|
||||
"""JSON doesn't distinguish int/float — int 1 is a valid threshold."""
|
||||
cfg = VideoStreamConfig.from_dict({"evs_threshold": 1})
|
||||
assert cfg.evs_threshold == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Frame buffer & sliding window (uses deque)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestFrameBuffer:
|
||||
def test_add_frame_basic(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, max_frames=10)
|
||||
session = VideoStreamSession(cfg)
|
||||
assert session.add_frame(make_jpeg()) is True
|
||||
assert session.frame_count == 1
|
||||
|
||||
def test_add_frame_too_large(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, max_frames=10)
|
||||
session = VideoStreamSession(cfg)
|
||||
huge = b"\xff" * (10 * 1024 * 1024 + 1) # just over 10 MB
|
||||
with pytest.raises(ValueError, match="Frame too large"):
|
||||
session.add_frame(huge)
|
||||
|
||||
def test_sliding_window(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, max_frames=3)
|
||||
session = VideoStreamSession(cfg)
|
||||
for i in range(5):
|
||||
session.add_frame(make_jpeg(r=i * 50))
|
||||
assert session.frame_count == 3
|
||||
|
||||
def test_sliding_window_keeps_newest(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, max_frames=2)
|
||||
session = VideoStreamSession(cfg)
|
||||
f1 = make_jpeg(10, 10, 10)
|
||||
f2 = make_jpeg(20, 20, 20)
|
||||
f3 = make_jpeg(30, 30, 30)
|
||||
session.add_frame(f1)
|
||||
session.add_frame(f2)
|
||||
session.add_frame(f3)
|
||||
sampled = session.sample_frames()
|
||||
assert len(sampled) == 2
|
||||
assert sampled[0] == f2
|
||||
assert sampled[1] == f3
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# EVS integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEVSIntegration:
|
||||
def test_evs_drops_identical_frames(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.90)
|
||||
session = VideoStreamSession(cfg)
|
||||
frame = make_jpeg(100, 100, 100)
|
||||
assert session.add_frame(frame) is True
|
||||
assert session.add_frame(frame) is False
|
||||
assert session.frame_count == 1
|
||||
|
||||
def test_evs_keeps_different_frames(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.95)
|
||||
session = VideoStreamSession(cfg)
|
||||
for i in range(5):
|
||||
assert session.add_frame(make_gradient_jpeg(seed=i)) is True
|
||||
assert session.frame_count == 5
|
||||
|
||||
def test_evs_disabled(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False)
|
||||
session = VideoStreamSession(cfg)
|
||||
frame = make_jpeg()
|
||||
assert session.add_frame(frame) is True
|
||||
assert session.add_frame(frame) is True
|
||||
assert session.frame_count == 2
|
||||
|
||||
def test_evs_stats(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.90)
|
||||
session = VideoStreamSession(cfg)
|
||||
frame = make_jpeg()
|
||||
session.add_frame(frame)
|
||||
session.add_frame(frame)
|
||||
stats = session.evs_stats
|
||||
assert stats is not None
|
||||
assert stats["retained_count"] == 1
|
||||
assert stats["dropped_count"] == 1
|
||||
|
||||
def test_evs_stats_none_when_disabled(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False)
|
||||
session = VideoStreamSession(cfg)
|
||||
assert session.evs_stats is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Uniform sampling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSampling:
|
||||
def test_sample_exact(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=3, max_frames=10)
|
||||
session = VideoStreamSession(cfg)
|
||||
frames = [make_gradient_jpeg(i) for i in range(3)]
|
||||
for f in frames:
|
||||
session.add_frame(f)
|
||||
sampled = session.sample_frames()
|
||||
assert len(sampled) == 3
|
||||
assert sampled == frames
|
||||
|
||||
def test_sample_fewer_than_requested(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=10, max_frames=64)
|
||||
session = VideoStreamSession(cfg)
|
||||
for i in range(3):
|
||||
session.add_frame(make_gradient_jpeg(i))
|
||||
assert len(session.sample_frames()) == 3
|
||||
|
||||
def test_sample_uniform(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=4, max_frames=64)
|
||||
session = VideoStreamSession(cfg)
|
||||
for i in range(10):
|
||||
session.add_frame(make_gradient_jpeg(i))
|
||||
sampled = session.sample_frames()
|
||||
assert len(sampled) == 4
|
||||
expected = [session._frames[i] for i in [0, 3, 6, 9]]
|
||||
assert sampled == expected
|
||||
|
||||
def test_sample_empty(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False)
|
||||
session = VideoStreamSession(cfg)
|
||||
assert session.sample_frames() == []
|
||||
|
||||
def test_sample_single_frame_from_multi_frame_buffer(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=1, max_frames=64)
|
||||
session = VideoStreamSession(cfg)
|
||||
first = make_gradient_jpeg(0)
|
||||
second = make_gradient_jpeg(1)
|
||||
session.add_frame(first)
|
||||
session.add_frame(second)
|
||||
|
||||
assert session.sample_frames() == [second]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Audio buffer (Phase 3)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestAudioBuffer:
|
||||
def test_add_audio_chunk(self):
|
||||
session = VideoStreamSession(VideoStreamConfig())
|
||||
assert session.has_audio is False
|
||||
session.add_audio_chunk(b"\x00" * 100)
|
||||
assert session.has_audio is True
|
||||
|
||||
def test_clear_audio(self):
|
||||
session = VideoStreamSession(VideoStreamConfig())
|
||||
session.add_audio_chunk(b"\x00" * 100)
|
||||
session.clear_audio()
|
||||
assert session.has_audio is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# build_chat_request
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBuildChatRequest:
|
||||
def test_video_only_request(self):
|
||||
cfg = VideoStreamConfig(model="test-model", evs_enabled=False, num_sample_frames=4)
|
||||
session = VideoStreamSession(cfg)
|
||||
for i in range(4):
|
||||
session.add_frame(make_gradient_jpeg(i))
|
||||
|
||||
request = session.build_chat_request("Describe this scene.")
|
||||
assert request.model == "test-model"
|
||||
assert request.stream is True
|
||||
|
||||
content = request.messages[0]["content"]
|
||||
assert len(content) == 5 # 4 image_url + 1 text
|
||||
image_parts = [p for p in content if p["type"] == "image_url"]
|
||||
text_parts = [p for p in content if p["type"] == "text"]
|
||||
assert len(image_parts) == 4
|
||||
assert len(text_parts) == 1
|
||||
assert text_parts[0]["text"] == "Describe this scene."
|
||||
|
||||
mm_kw = getattr(request, "mm_processor_kwargs", None)
|
||||
assert mm_kw is None or not mm_kw.get("use_audio_in_video", False)
|
||||
|
||||
def test_video_plus_audio_request(self):
|
||||
cfg = VideoStreamConfig(model="test-model", evs_enabled=False, num_sample_frames=2)
|
||||
session = VideoStreamSession(cfg)
|
||||
session.add_frame(make_gradient_jpeg(0))
|
||||
session.add_frame(make_gradient_jpeg(1))
|
||||
session.add_audio_chunk(b"\x00" * 3200)
|
||||
|
||||
request = session.build_chat_request("What is being said?")
|
||||
|
||||
content = request.messages[0]["content"]
|
||||
assert len(content) == 4 # 2 image_url + 1 audio_url + 1 text
|
||||
audio_parts = [p for p in content if p["type"] == "audio_url"]
|
||||
assert len(audio_parts) == 1
|
||||
# RFC 3551: audio/L16 for linear 16-bit PCM
|
||||
assert audio_parts[0]["audio_url"]["url"].startswith("data:audio/L16;rate=16000;base64,")
|
||||
|
||||
mm_kw = getattr(request, "mm_processor_kwargs", None) or {}
|
||||
assert mm_kw.get("use_audio_in_video") is True
|
||||
|
||||
def test_image_url_is_valid_base64(self):
|
||||
cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=1)
|
||||
session = VideoStreamSession(cfg)
|
||||
session.add_frame(make_jpeg(200, 100, 50))
|
||||
request = session.build_chat_request("test")
|
||||
content = request.messages[0]["content"]
|
||||
img_url = content[0]["image_url"]["url"]
|
||||
assert img_url.startswith("data:image/jpeg;base64,")
|
||||
import base64
|
||||
|
||||
b64_data = img_url.split(",", 1)[1]
|
||||
decoded = base64.b64decode(b64_data)
|
||||
assert len(decoded) > 0
|
||||
|
||||
def test_clear_audio_after_query(self):
|
||||
session = VideoStreamSession(VideoStreamConfig(evs_enabled=False))
|
||||
session.add_frame(make_jpeg())
|
||||
session.add_audio_chunk(b"\x00" * 100)
|
||||
session.build_chat_request("test")
|
||||
session.clear_audio()
|
||||
assert session.has_audio is False
|
||||
assert session.frame_count == 1
|
||||
Reference in new issue
Block a user