zcbot/tests/test_attachments.py

78 lines
3.0 KiB
Python

import tempfile
import unittest
from pathlib import Path
from core.artifacts import ArtifactPathError
from core.attachments import (
content_for_model,
materialize_native_images,
normalize_attachment_refs,
)
class AttachmentRefTests(unittest.TestCase):
def test_normalizes_user_relative_path_to_task_relative_metadata(self):
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
wd = root / "实验分析"
wd.mkdir()
image = wd / "显微照片.png"
image.write_bytes(b"png")
refs = normalize_attachment_refs(
[{"path": "实验分析/显微照片.png", "kind": "image"}],
working_dir=wd,
user_root=root,
)
self.assertEqual(refs[0]["path"], "显微照片.png")
self.assertEqual(refs[0]["kind"], "image")
self.assertEqual(refs[0]["size_bytes"], 3)
self.assertEqual(refs[0]["scope"], "working_dir")
def test_rejects_attachment_outside_working_dir(self):
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
wd = root / "任务"
wd.mkdir()
(root / "outside.txt").write_text("x", encoding="utf-8")
with self.assertRaises(ArtifactPathError):
normalize_attachment_refs(
[{"path": "../outside.txt", "kind": "file"}],
working_dir=wd,
user_root=root,
)
def test_builds_provider_content_without_changing_stored_text(self):
refs = [{"path": "figures/a.png", "kind": "image"}]
self.assertEqual(
content_for_model("请分析", refs, "实验分析"),
"请分析\n\n[用户上传的参考图] 实验分析/figures/a.png",
)
self.assertEqual(
content_for_model("", refs, "实验分析"),
"[用户上传的参考图] 实验分析/figures/a.png",
)
def test_materializes_native_image_without_persisting_internal_fields(self):
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
wd = root / "任务"
wd.mkdir()
(wd / "a.png").write_bytes(b"png")
messages = [{
"role": "user",
"content": "请分析\n\n[用户上传的参考图] 任务/a.png",
"_attachment_refs": [{"path": "a.png", "kind": "image"}],
}]
prepared, count = materialize_native_images(
messages, enabled=True, working_dir=wd, user_root=root,
)
self.assertEqual(count, 1)
self.assertNotIn("_attachment_refs", prepared[0])
self.assertEqual(prepared[0]["content"][0]["type"], "text")
self.assertEqual(prepared[0]["content"][1]["type"], "image_url")
self.assertTrue(
prepared[0]["content"][1]["image_url"]["url"].startswith(
"data:image/png;base64,"
)
)