75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
|
|
||
|
|
import base64
|
||
|
|
import io
|
||
|
|
import zipfile
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
from fastapi.responses import StreamingResponse
|
||
|
|
|
||
|
|
from vllm_omni.entrypoints.openai.protocol.images import (
|
||
|
|
_FILE_RESPONSE_CHUNK_SIZE,
|
||
|
|
ImageData,
|
||
|
|
ImageGenerationResponse,
|
||
|
|
)
|
||
|
|
|
||
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||
|
|
|
||
|
|
_NEWLINE_HEAVY_PAYLOAD = (b"\x89PNG\n\x1a\n" * 64) + bytes(range(256)) * 8
|
||
|
|
|
||
|
|
|
||
|
|
def _response(*payloads: bytes) -> ImageGenerationResponse:
|
||
|
|
return ImageGenerationResponse(
|
||
|
|
created=0,
|
||
|
|
data=[ImageData(b64_json=base64.b64encode(payload).decode()) for payload in payloads],
|
||
|
|
output_format="png",
|
||
|
|
size="1x1",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def _collect(response: StreamingResponse) -> list[bytes]:
|
||
|
|
return [bytes(chunk) async for chunk in response.body_iterator]
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_single_file_response_chunks_ignore_newlines() -> None:
|
||
|
|
payload = _NEWLINE_HEAVY_PAYLOAD
|
||
|
|
assert b"\n" in payload
|
||
|
|
|
||
|
|
response = _response(payload).stream_response()
|
||
|
|
chunks = await _collect(response)
|
||
|
|
|
||
|
|
assert len(chunks) == 1
|
||
|
|
assert b"".join(chunks) == payload
|
||
|
|
assert response.headers["content-length"] == str(len(payload))
|
||
|
|
assert response.headers["content-type"] == "image/png"
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_single_file_response_chunks_are_bounded() -> None:
|
||
|
|
payload = _NEWLINE_HEAVY_PAYLOAD * 200
|
||
|
|
expected_chunks = -(-len(payload) // _FILE_RESPONSE_CHUNK_SIZE)
|
||
|
|
assert expected_chunks > 1
|
||
|
|
|
||
|
|
chunks = await _collect(_response(payload).stream_response())
|
||
|
|
|
||
|
|
assert len(chunks) == expected_chunks
|
||
|
|
assert max(len(chunk) for chunk in chunks) <= _FILE_RESPONSE_CHUNK_SIZE
|
||
|
|
assert b"".join(chunks) == payload
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.asyncio
|
||
|
|
async def test_zip_file_response_streams_whole_archive() -> None:
|
||
|
|
payloads = (_NEWLINE_HEAVY_PAYLOAD, b"second image\nwith a newline")
|
||
|
|
|
||
|
|
response = _response(*payloads).stream_response()
|
||
|
|
chunks = await _collect(response)
|
||
|
|
|
||
|
|
zip_bytes = b"".join(chunks)
|
||
|
|
assert len(chunks) == -(-len(zip_bytes) // _FILE_RESPONSE_CHUNK_SIZE)
|
||
|
|
assert response.headers["content-length"] == str(len(zip_bytes))
|
||
|
|
assert response.headers["content-type"] == "application/zip"
|
||
|
|
with zipfile.ZipFile(io.BytesIO(zip_bytes)) as archive:
|
||
|
|
assert archive.read("image_0.png") == payloads[0]
|
||
|
|
assert archive.read("image_1.png") == payloads[1]
|