119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
|
from io import BytesIO
|
||
|
|
import unittest
|
||
|
|
from unittest.mock import AsyncMock, patch
|
||
|
|
|
||
|
|
from PIL import Image
|
||
|
|
from pydantic import ValidationError
|
||
|
|
|
||
|
|
from models import AssistantConfig
|
||
|
|
from services.test_assets import read_test_image, store_test_image
|
||
|
|
from services.test_runs.text_runner import _transcript_content
|
||
|
|
from services.vision import build_image_user_message
|
||
|
|
from test_schemas import FixedInputTurn, TestImageRef
|
||
|
|
|
||
|
|
|
||
|
|
def _jpeg_bytes() -> bytes:
|
||
|
|
buffer = BytesIO()
|
||
|
|
Image.new("RGB", (12, 8), (220, 40, 40)).save(buffer, format="JPEG")
|
||
|
|
return buffer.getvalue()
|
||
|
|
|
||
|
|
|
||
|
|
def _image_ref() -> TestImageRef:
|
||
|
|
return TestImageRef(
|
||
|
|
asset_id="ta_fixture",
|
||
|
|
filename="fixture.jpg",
|
||
|
|
width=12,
|
||
|
|
height=8,
|
||
|
|
size_bytes=128,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class TestImageSchemaTests(unittest.TestCase):
|
||
|
|
def test_allows_image_only_turn(self):
|
||
|
|
turn = FixedInputTurn(id="turn_1", image=_image_ref())
|
||
|
|
|
||
|
|
self.assertEqual(turn.user_input, "")
|
||
|
|
self.assertEqual(turn.image.asset_id, "ta_fixture")
|
||
|
|
|
||
|
|
def test_rejects_turn_without_text_or_image(self):
|
||
|
|
with self.assertRaisesRegex(ValidationError, "至少填写一项"):
|
||
|
|
FixedInputTurn(id="turn_1")
|
||
|
|
|
||
|
|
|
||
|
|
class PersistentTestAssetTests(unittest.TestCase):
|
||
|
|
def test_normalizes_and_stores_under_test_prefix(self):
|
||
|
|
with patch("services.test_assets.put_object") as put_object:
|
||
|
|
stored = store_test_image("ta_fixture", _jpeg_bytes())
|
||
|
|
|
||
|
|
self.assertEqual(stored.storage_key, "test-assets/ta_fixture.jpg")
|
||
|
|
self.assertEqual((stored.width, stored.height), (12, 8))
|
||
|
|
self.assertEqual(put_object.call_args.args[0], stored.storage_key)
|
||
|
|
self.assertEqual(put_object.call_args.args[2], "image/jpeg")
|
||
|
|
|
||
|
|
def test_reads_only_test_asset_keys(self):
|
||
|
|
with patch("services.test_assets.get_object", return_value=b"jpeg"):
|
||
|
|
self.assertEqual(read_test_image("test-assets/ta_fixture.jpg"), b"jpeg")
|
||
|
|
with self.assertRaisesRegex(ValueError, "存储位置无效"):
|
||
|
|
read_test_image("conversation-inputs/unsafe.jpg")
|
||
|
|
|
||
|
|
|
||
|
|
class ImageUserMessageTests(unittest.IsolatedAsyncioTestCase):
|
||
|
|
async def test_native_model_receives_multimodal_message(self):
|
||
|
|
message = await build_image_user_message(
|
||
|
|
AssistantConfig(
|
||
|
|
type="prompt",
|
||
|
|
vision_enabled=True,
|
||
|
|
llm_support_image_input=True,
|
||
|
|
),
|
||
|
|
_jpeg_bytes(),
|
||
|
|
"这是什么颜色？",
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(message["role"], "user")
|
||
|
|
self.assertEqual(message["content"][0]["text"], "这是什么颜色？")
|
||
|
|
self.assertTrue(
|
||
|
|
message["content"][1]["image_url"]["url"].startswith(
|
||
|
|
"data:image/jpeg;base64,"
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
async def test_auxiliary_model_observation_becomes_text_context(self):
|
||
|
|
with patch(
|
||
|
|
"services.vision.analyze_image_with_vision_model",
|
||
|
|
new=AsyncMock(return_value="画面中是一个红色方块。"),
|
||
|
|
):
|
||
|
|
message = await build_image_user_message(
|
||
|
|
AssistantConfig(
|
||
|
|
type="prompt",
|
||
|
|
vision_enabled=True,
|
||
|
|
vision_model_resource_id="vision_model",
|
||
|
|
vision_llm_support_image_input=True,
|
||
|
|
),
|
||
|
|
_jpeg_bytes(),
|
||
|
|
"画面里有什么？",
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertIsInstance(message["content"], str)
|
||
|
|
self.assertIn("红色方块", message["content"])
|
||
|
|
|
||
|
|
def test_transcript_does_not_persist_image_data(self):
|
||
|
|
content = _transcript_content(
|
||
|
|
{
|
||
|
|
"role": "user",
|
||
|
|
"content": [
|
||
|
|
{"type": "text", "text": "请看图"},
|
||
|
|
{
|
||
|
|
"type": "image_url",
|
||
|
|
"image_url": {"url": "data:image/jpeg;base64,secret"},
|
||
|
|
},
|
||
|
|
],
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
self.assertEqual(content, "请看图\n[用户提交了一张测试图片]")
|
||
|
|
self.assertNotIn("base64", content)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|