86 lines
3.3 KiB
Python
86 lines
3.3 KiB
Python
import base64
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from video_ai_analysis_poc.vlm_client import infer_clip
|
|
|
|
|
|
class VlmClientTests(unittest.TestCase):
|
|
def test_infer_clip_uses_config_prompt_url_and_data_uri_images(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
output_dir = Path(tmp)
|
|
frame_path = output_dir / "frames" / "video-abc" / "000001.jpg"
|
|
frame_path.parent.mkdir(parents=True)
|
|
frame_path.write_bytes(b"jpg-bytes")
|
|
calls = []
|
|
|
|
def http_post(url, payload, timeout_seconds):
|
|
calls.append((url, payload, timeout_seconds))
|
|
return {
|
|
"status": 200,
|
|
"body": {
|
|
"choices": [
|
|
{
|
|
"message": {
|
|
"content": json.dumps(
|
|
{"screen_time": "10:00:01", "events": []}
|
|
)
|
|
}
|
|
}
|
|
]
|
|
},
|
|
}
|
|
|
|
result = infer_clip(
|
|
{
|
|
"clip_id": "video-abc_c000001",
|
|
"frame_times": [
|
|
{
|
|
"frame_path": "frames/video-abc/000001.jpg",
|
|
"offset_seconds": 0.0,
|
|
"timecode": "00:00:00",
|
|
}
|
|
],
|
|
},
|
|
output_dir,
|
|
{
|
|
"api_base_url": "http://localhost:8679/",
|
|
"chat_completions_path": "/v1/chat/completions",
|
|
"model": "memai-zhengxin-v3-20260413",
|
|
"timeout_seconds": 17,
|
|
"max_tokens": 256,
|
|
"temperature": 0,
|
|
"image_transport": "data_uri",
|
|
},
|
|
{
|
|
"system": "system prompt from config",
|
|
"user": "user prompt from config",
|
|
},
|
|
http_post=http_post,
|
|
)
|
|
|
|
self.assertEqual(result["raw_response"], '{"screen_time": "10:00:01", "events": []}')
|
|
self.assertEqual(len(calls), 1)
|
|
url, payload, timeout_seconds = calls[0]
|
|
self.assertEqual(url, "http://localhost:8679/v1/chat/completions")
|
|
self.assertEqual(timeout_seconds, 17)
|
|
self.assertEqual(payload["model"], "memai-zhengxin-v3-20260413")
|
|
self.assertEqual(payload["messages"][0]["role"], "system")
|
|
self.assertEqual(payload["messages"][0]["content"], "system prompt from config")
|
|
user_content = payload["messages"][1]["content"]
|
|
self.assertEqual(user_content[0], {"type": "text", "text": "user prompt from config"})
|
|
self.assertEqual(user_content[1]["type"], "image_url")
|
|
expected_data = base64.b64encode(b"jpg-bytes").decode("ascii")
|
|
self.assertEqual(
|
|
user_content[1]["image_url"]["url"],
|
|
f"data:image/jpeg;base64,{expected_data}",
|
|
)
|
|
self.assertEqual(result["http_status"], 200)
|
|
self.assertIsInstance(result["latency_ms"], int)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|