zl-q
7b8865e256
- 新增 stepStarted/stepFinished 事件类型支持 - 前端实现图片附件上传和预览功能 - 后端增强工具结果存储和事件处理 - 完善相关单元测试和集成测试
86 lines
2.4 KiB
Python
86 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
import v1.agent.attachment_storage as attachment_storage_module
|
|
|
|
|
|
class _FakeBucket:
|
|
def __init__(self) -> None:
|
|
self.upload_calls: list[tuple[str, bytes, dict[str, str]]] = []
|
|
self.download_calls: list[str] = []
|
|
|
|
def upload(self, path: str, content: bytes, options: dict[str, str]) -> object:
|
|
self.upload_calls.append((path, content, options))
|
|
return {"path": path}
|
|
|
|
def download(self, path: str) -> object:
|
|
self.download_calls.append(path)
|
|
return b"ok"
|
|
|
|
|
|
class _FakeStorage:
|
|
def __init__(self, bucket: _FakeBucket) -> None:
|
|
self._bucket = bucket
|
|
|
|
def from_(self, bucket: str) -> object:
|
|
del bucket
|
|
return self._bucket
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_attachment_storage_rejects_unexpected_bucket(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
storage = attachment_storage_module.AgentAttachmentStorage()
|
|
monkeypatch.setattr(
|
|
attachment_storage_module.config.storage,
|
|
"bucket",
|
|
"allowed-bucket",
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="Invalid attachment bucket"):
|
|
await storage.upload_bytes(
|
|
bucket="other-bucket",
|
|
path="agent-inputs/u/t/r/file.png",
|
|
content=b"data",
|
|
content_type="image/png",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_attachment_storage_accepts_configured_bucket(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
storage = attachment_storage_module.AgentAttachmentStorage()
|
|
fake_bucket = _FakeBucket()
|
|
fake_client = SimpleNamespace(storage=_FakeStorage(fake_bucket))
|
|
monkeypatch.setattr(
|
|
attachment_storage_module.config.storage,
|
|
"bucket",
|
|
"allowed-bucket",
|
|
)
|
|
monkeypatch.setattr(
|
|
attachment_storage_module.supabase_service,
|
|
"get_admin_client",
|
|
lambda: fake_client,
|
|
)
|
|
|
|
path = await storage.upload_bytes(
|
|
bucket="allowed-bucket",
|
|
path="agent-inputs/u/t/r/file.png",
|
|
content=b"data",
|
|
content_type="image/png",
|
|
)
|
|
payload = await storage.download_bytes(
|
|
bucket="allowed-bucket",
|
|
path=path,
|
|
)
|
|
|
|
assert path == "agent-inputs/u/t/r/file.png"
|
|
assert payload == b"ok"
|
|
assert len(fake_bucket.upload_calls) == 1
|
|
assert fake_bucket.download_calls == ["agent-inputs/u/t/r/file.png"]
|