vllm-project--vllm-omni
wehub-resource-sync
596 行
21 KiB
Python
596 行
21 KiB
Python
# 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 ErrorInfo, ErrorResponse
|
|
|
|
err = ErrorResponse(
|
|
error=ErrorInfo(
|
|
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("你好🌍")) == ["你好🌍"]
|