| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973 |
- from __future__ import annotations
- import subprocess
- import sys
- import tempfile
- import time
- import unittest
- from concurrent.futures import ThreadPoolExecutor
- from io import BytesIO
- from pathlib import Path
- from types import SimpleNamespace
- from typing import Any
- from unittest.mock import MagicMock, patch
- import requests
- from PIL import Image
- from production_build_agents.run.operation_journal import (
- OperationOutcomeUnknownError,
- UncertainSideEffectError,
- )
- from production_build_agents.run.records import (
- VersionConflictError,
- )
- from production_build_agents.tools.discovery import RemoteToolClient
- from production_build_agents.tools.media import (
- concat_videos,
- crop_image,
- extract_frames,
- fit_audio_duration,
- grid_collage,
- image_as_data_url,
- mix_audio_tracks,
- mux_audio,
- overlay_text,
- probe_media,
- trim_audio,
- trim_video,
- )
- from production_build_agents.tools.publishing import publish_file
- from production_build_agents.tools.registry import (
- ToolRegistryError,
- create_default_tool_registry,
- )
- class FakeRemoteToolClient:
- def search(self, query: str, *, limit: int = 5) -> dict[str, Any]:
- return {
- "success": True,
- "results": [
- {"tool_id": "fake_generate", "summary": query},
- {"tool_id": "other_generate", "summary": "未授权工具"},
- ],
- "total": 2,
- }
- def inspect(self, tool_ids: list[str]) -> dict[str, Any]:
- return {
- "success": True,
- "tools": {
- tool_id: {"success": True, "input_schema": {}}
- for tool_id in tool_ids
- },
- }
- def invoke(self, tool_id: str, payload: dict[str, Any]) -> dict[str, Any]:
- return {
- "success": True,
- "data": {"tool_id": tool_id, "payload": payload},
- }
- class FakeResponse:
- def __init__(self, payload: dict[str, Any]) -> None:
- self.payload = payload
- def raise_for_status(self) -> None:
- return None
- def json(self) -> dict[str, Any]:
- return self.payload
- class FakeSession:
- def __init__(self) -> None:
- self.requests: list[tuple[str, str]] = []
- self.invoke_attempts = 0
- def post(self, url: str, **_: Any) -> FakeResponse:
- self.requests.append(("POST", url))
- return FakeResponse(
- {
- "data": {
- "items": [
- {
- "tool_id": "fake_generate",
- "description": "测试生成工具",
- }
- ]
- }
- }
- )
- def get(self, url: str, **_: Any) -> FakeResponse:
- self.requests.append(("GET", url))
- return FakeResponse(
- {
- "data": {
- "invoke": {
- "base_url": "https://invoke.test/fake_generate",
- "request_method": "POST",
- "input_schema": {"type": "object"},
- }
- }
- }
- )
- def request(self, method: str, url: str, **_: Any) -> FakeResponse:
- self.requests.append((method, url))
- self.invoke_attempts += 1
- if self.invoke_attempts == 1:
- raise requests.ConnectionError("temporary connection failure")
- return FakeResponse({"code": 0, "data": {"url": "https://result.test/a"}})
- class ToolRegistryTest(unittest.TestCase):
- def test_formal_seedance_tool_localizes_and_silences_result(self) -> None:
- class SeedanceClient:
- def invoke(self, tool_id: str, payload: dict[str, Any]) -> dict[str, Any]:
- self.tool_id = tool_id
- self.payload = payload
- return {
- "success": True,
- "outcome": "SUCCEEDED",
- "data": {"result_url": "https://result.test/shot.mp4"},
- }
- client = SeedanceClient()
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- downloaded = root / "downloaded.mp4"
- silent = root / "silent.mp4"
- downloaded.touch()
- silent.touch()
- with (
- patch(
- "production_build_agents.tools.video_tools."
- "stable_fetch_media",
- return_value=downloaded,
- ),
- patch(
- "production_build_agents.tools.video_tools.probe_media",
- side_effect=[
- {"media_type": "video", "has_audio": True},
- {
- "media_type": "video",
- "has_audio": False,
- "duration_sec": 4.0,
- "width": 1080,
- "height": 1920,
- "codec": "h264",
- },
- ],
- ),
- patch(
- "production_build_agents.tools.video_tools."
- "strip_video_audio",
- return_value={"local_paths": [str(silent)]},
- ),
- ):
- tool = create_default_tool_registry(
- output_dir=root,
- remote_client=client,
- ).resolve(("generate_seedance_video",))[0]
- result = tool.invoke(
- {
- "prompt": "镜头缓慢推进",
- "first_frame_url": "https://published.test/anchor.png",
- "duration": 4,
- }
- )
- self.assertTrue(result["success"])
- self.assertEqual(client.tool_id, "seedance_generate_video")
- self.assertFalse(client.payload["generate_audio"])
- self.assertFalse(client.payload["web_search"])
- self.assertEqual(result["local_path"], str(silent.resolve()))
- self.assertFalse(result["has_audio"])
- def test_seedance_localization_failure_is_sealed_as_unknown(self) -> None:
- client = MagicMock()
- client.invoke.return_value = {
- "success": True,
- "outcome": "SUCCEEDED",
- "data": {"result_url": "https://result.test/shot.mp4"},
- }
- with tempfile.TemporaryDirectory() as temp_dir, patch(
- "production_build_agents.tools.video_tools.stable_fetch_media",
- side_effect=OSError("download interrupted"),
- ):
- root = Path(temp_dir)
- tool = create_default_tool_registry(
- output_dir=root / "outputs",
- remote_client=client,
- run_dir=root,
- executor_run_id="seedance-executor",
- ).resolve(("generate_seedance_video",))[0]
- params = {
- "prompt": "镜头缓慢推进",
- "first_frame_url": "https://published.test/anchor.png",
- "duration": 4,
- }
- with self.assertRaises(OperationOutcomeUnknownError):
- tool.invoke(params)
- client.invoke.assert_called_once()
- def test_remote_media_is_downloaded_once_per_run(self) -> None:
- image_buffer = BytesIO()
- Image.new("RGB", (8, 8), "navy").save(image_buffer, format="PNG")
- response = MagicMock()
- response.__enter__.return_value = response
- response.headers = {"content-type": "image/png"}
- response.iter_content.return_value = [image_buffer.getvalue()]
- with tempfile.TemporaryDirectory() as temp_dir, patch(
- "production_build_agents.tools.media._SESSION.get",
- return_value=response,
- ) as get:
- root = Path(temp_dir)
- source = "https://example.test/reference"
- probe_media(source, output_dir=root)
- data_url = image_as_data_url(source, root)
- self.assertTrue(data_url.startswith("data:image/png;base64,"))
- self.assertEqual(get.call_count, 1)
- def test_probe_media_identifies_image_video_and_audio(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- image = root / "reference.png"
- audio = root / "reference.wav"
- video = root / "reference.mp4"
- Image.new("RGB", (32, 18), "navy").save(image)
- subprocess.run(
- [
- "ffmpeg",
- "-y",
- "-f",
- "lavfi",
- "-i",
- "sine=frequency=440:duration=0.5",
- str(audio),
- ],
- check=True,
- capture_output=True,
- )
- subprocess.run(
- [
- "ffmpeg",
- "-y",
- "-f",
- "lavfi",
- "-i",
- "color=c=blue:s=64x36:d=0.5",
- "-pix_fmt",
- "yuv420p",
- str(video),
- ],
- check=True,
- capture_output=True,
- )
- image_info = probe_media(str(image), output_dir=root / "cache")
- audio_info = probe_media(str(audio), output_dir=root / "cache")
- video_info = probe_media(str(video), output_dir=root / "cache")
- self.assertEqual(image_info["media_type"], "image")
- self.assertEqual((image_info["width"], image_info["height"]), (32, 18))
- self.assertEqual(audio_info["media_type"], "audio")
- self.assertEqual(audio_info["channels"], 1)
- self.assertEqual(video_info["media_type"], "video")
- self.assertEqual((video_info["width"], video_info["height"]), (64, 36))
- def test_view_images_keeps_base64_out_of_tool_message_payload(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- image_path = root / "source.png"
- Image.new("RGB", (8, 8), "blue").save(image_path)
- view_images = create_default_tool_registry(
- output_dir=root / "outputs",
- remote_client=FakeRemoteToolClient(),
- ).resolve(("view_images",))[0]
- result = view_images.invoke(
- {"image_sources": [str(image_path)]}
- )
- self.assertTrue(result["success"])
- self.assertEqual(result["images"][0]["source"], str(image_path))
- self.assertNotIn("data_url", result["images"][0])
- self.assertTrue(
- Path(result["images"][0]["local_path"]).is_absolute()
- )
- def test_complete_registry_and_skill_allowlist(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- registry = create_default_tool_registry(
- output_dir=Path(temp_dir),
- remote_client=FakeRemoteToolClient(),
- )
- self.assertEqual(
- set(registry.tool_ids),
- {
- "search_tool",
- "inspect_tool",
- "run_tool",
- "submit_reference_video",
- "fetch_generated_video",
- "generate_seedance_video",
- "probe_media",
- "detect_faces",
- "measure_audio_loudness",
- "view_images",
- "publish_media_reference",
- "crop_image",
- "grid_collage",
- "overlay_text",
- "extract_frames",
- "video_trim",
- "audio_trim",
- "audio_fit_duration",
- "mix_audio_tracks",
- "video_concat",
- "assemble_segment_media",
- "video_mux_audio",
- "synthesize_speech",
- "transcribe_audio",
- "create_ass_subtitles",
- "inspect_ass_subtitles",
- "render_ass_subtitles",
- },
- )
- self.assertEqual(
- [tool.name for tool in registry.resolve(("view_images", "run_tool"))],
- ["view_images", "run_tool"],
- )
- with self.assertRaisesRegex(ToolRegistryError, "未注册工具"):
- registry.resolve(("not_exists",))
- def test_tools_readme_covers_every_registered_tool(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- registry = create_default_tool_registry(
- output_dir=Path(temp_dir),
- remote_client=FakeRemoteToolClient(),
- )
- readme = (
- Path(__file__).parents[2]
- / "production_build_agents"
- / "tools"
- / "README.md"
- ).read_text(encoding="utf-8")
- for tool_id, registered in registry.tools.items():
- self.assertIn(f"`{tool_id}`", readme)
- self.assertGreater(len(registered.description), 60)
- def test_remote_discovery_inspection_and_invocation(self) -> None:
- class SlowRemoteToolClient(FakeRemoteToolClient):
- def search(
- self,
- query: str,
- *,
- limit: int = 5,
- ) -> dict[str, Any]:
- time.sleep(0.01)
- return super().search(query, limit=limit)
- with tempfile.TemporaryDirectory() as temp_dir:
- registry = create_default_tool_registry(
- output_dir=Path(temp_dir),
- remote_client=SlowRemoteToolClient(),
- allowed_remote_tool_ids={"fake_generate"},
- )
- search, inspect, run = registry.resolve(
- ("search_tool", "inspect_tool", "run_tool")
- )
- searched = search.invoke({"query": "生成图片"})
- self.assertTrue(searched["success"])
- self.assertEqual(
- [item["tool_id"] for item in searched["results"]],
- ["fake_generate"],
- )
- timed = search.invoke({"query": "生成图片"})
- self.assertGreaterEqual(timed["_duration_ms"], 5)
- self.assertTrue(
- inspect.invoke({"tool_ids": ["fake_generate"]})["success"]
- )
- denied_inspect = inspect.invoke(
- {"tool_ids": ["other_generate"]}
- )
- self.assertFalse(denied_inspect["success"])
- result = run.invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "test"},
- }
- )
- self.assertEqual(result["data"]["tool_id"], "fake_generate")
- denied = run.invoke(
- {"tool_id": "other_generate", "params": {}}
- )
- self.assertFalse(denied["success"])
- def test_run_tool_is_blocked_before_remote_call_without_inspect(
- self,
- ) -> None:
- class CountingClient(FakeRemoteToolClient):
- def __init__(self) -> None:
- self.calls = 0
- def invoke(self, tool_id: str, payload: dict[str, Any]) -> dict:
- self.calls += 1
- return super().invoke(tool_id, payload)
- client = CountingClient()
- with tempfile.TemporaryDirectory() as temp_dir:
- run_tool = create_default_tool_registry(
- output_dir=Path(temp_dir),
- remote_client=client,
- allowed_remote_tool_ids={"fake_generate"},
- ).resolve(("run_tool",))[0]
- result = run_tool.invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "不得发送"},
- }
- )
- self.assertFalse(result["success"])
- self.assertIn("必须成功 inspect", result["error"])
- self.assertEqual(client.calls, 0)
- def test_run_tool_enforces_and_restores_remote_call_budget(
- self,
- ) -> None:
- class CountingClient(FakeRemoteToolClient):
- def __init__(self) -> None:
- self.calls = 0
- def invoke(
- self,
- tool_id: str,
- payload: dict[str, Any],
- ) -> dict[str, Any]:
- self.calls += 1
- return super().invoke(tool_id, payload)
- client = CountingClient()
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- registry = create_default_tool_registry(
- output_dir=root / "outputs",
- remote_client=client,
- allowed_remote_tool_ids={"fake_generate"},
- inspected_remote_tool_ids={"fake_generate"},
- remote_tool_call_limits={"fake_generate": 2},
- )
- run_tool = registry.resolve(("run_tool",))[0]
- for index in range(2):
- self.assertTrue(
- run_tool.invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": str(index)},
- }
- )["success"]
- )
- denied = run_tool.invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "over-budget"},
- }
- )
- restored = create_default_tool_registry(
- output_dir=root / "restored-outputs",
- remote_client=client,
- allowed_remote_tool_ids={"fake_generate"},
- inspected_remote_tool_ids={"fake_generate"},
- remote_tool_call_limits={"fake_generate": 2},
- remote_tool_call_counts={"fake_generate": 2},
- ).resolve(("run_tool",))[0]
- denied_after_restore = restored.invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "after-restore"},
- }
- )
- self.assertFalse(denied["success"])
- self.assertFalse(denied_after_restore["success"])
- self.assertIn("调用上限 2", denied["error"])
- self.assertIn("调用上限 2", denied_after_restore["error"])
- self.assertEqual(client.calls, 2)
- def test_run_tool_budget_is_atomic_for_parallel_calls(self) -> None:
- class CountingClient(FakeRemoteToolClient):
- def __init__(self) -> None:
- self.calls: list[dict[str, Any]] = []
- def invoke(
- self,
- tool_id: str,
- payload: dict[str, Any],
- ) -> dict[str, Any]:
- self.calls.append(payload)
- time.sleep(0.01)
- return super().invoke(tool_id, payload)
- client = CountingClient()
- with tempfile.TemporaryDirectory() as temp_dir:
- run_tool = create_default_tool_registry(
- output_dir=Path(temp_dir),
- remote_client=client,
- allowed_remote_tool_ids={"fake_generate"},
- inspected_remote_tool_ids={"fake_generate"},
- remote_tool_call_limits={"fake_generate": 2},
- ).resolve(("run_tool",))[0]
- with ThreadPoolExecutor(max_workers=8) as pool:
- results = list(
- pool.map(
- lambda index: run_tool.invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": str(index)},
- }
- ),
- range(8),
- )
- )
- self.assertEqual(
- sum(result["success"] is True for result in results),
- 2,
- )
- self.assertEqual(len(client.calls), 2)
- def test_remote_client_allows_anonymous_internal_service(self) -> None:
- with patch.dict(
- "os.environ",
- {
- "TOOLS_API_BASE_URL": "https://tools.test/v1/tools",
- "TOOLS_API_TOKEN": "",
- },
- ):
- client = RemoteToolClient.from_env()
- self.assertEqual(client.token, "")
- self.assertNotIn("Authorization", client._headers())
- self.assertEqual(
- client._headers(json_body=True),
- {"Content-Type": "application/json"},
- )
- def test_remote_client_never_retries_unknown_connection_failure(self) -> None:
- session = FakeSession()
- client = RemoteToolClient(
- base_url="https://tools.test/v1/tools",
- token="test-token",
- session=session,
- )
- self.assertTrue(client.search("生成图片")["success"])
- inspected = client.inspect(["fake_generate"])
- self.assertTrue(inspected["success"])
- result = client.invoke("fake_generate", {"prompt": "test"})
- self.assertFalse(result["success"])
- self.assertEqual(result["outcome"], "OUTCOME_UNKNOWN")
- self.assertEqual(
- [
- request
- for request in session.requests
- if request[1] == "https://invoke.test/fake_generate"
- ],
- [
- ("POST", "https://invoke.test/fake_generate"),
- ],
- )
- def test_remote_client_marks_success_and_known_failure(self) -> None:
- session = FakeSession()
- session.invoke_attempts = 1
- client = RemoteToolClient(
- base_url="https://tools.test/v1/tools",
- session=session,
- )
- result = client.invoke("fake_generate", {"prompt": "test"})
- self.assertTrue(result["success"])
- self.assertEqual(result["outcome"], "SUCCEEDED")
- def test_invalid_response_after_remote_post_is_unknown(self) -> None:
- class InvalidJsonResponse(FakeResponse):
- def json(self) -> dict[str, Any]:
- raise ValueError("invalid json")
- class InvalidJsonSession(FakeSession):
- def request(self, method: str, url: str, **_: Any):
- self.requests.append((method, url))
- return InvalidJsonResponse({})
- client = RemoteToolClient(
- base_url="https://tools.test/v1/tools",
- session=InvalidJsonSession(),
- )
- result = client.invoke("fake_generate", {"prompt": "test"})
- self.assertFalse(result["success"])
- self.assertEqual(result["outcome"], "OUTCOME_UNKNOWN")
- def test_http_408_after_remote_post_is_unknown(self) -> None:
- class TimeoutResponse(FakeResponse):
- status_code = 408
- def raise_for_status(self) -> None:
- raise requests.HTTPError(
- "request timeout",
- response=self,
- )
- class TimeoutSession(FakeSession):
- def request(self, method: str, url: str, **_: Any):
- self.requests.append((method, url))
- return TimeoutResponse({})
- client = RemoteToolClient(
- base_url="https://tools.test/v1/tools",
- session=TimeoutSession(),
- )
- result = client.invoke("fake_generate", {"prompt": "test"})
- self.assertFalse(result["success"])
- self.assertEqual(result["outcome"], "OUTCOME_UNKNOWN")
- def test_run_tool_accepts_more_than_five_reference_urls(self) -> None:
- client = FakeRemoteToolClient()
- with tempfile.TemporaryDirectory() as temp_dir:
- registry = create_default_tool_registry(
- output_dir=Path(temp_dir),
- remote_client=client,
- allowed_remote_tool_ids={"fake_generate"},
- inspected_remote_tool_ids={"fake_generate"},
- )
- run_tool = registry.resolve(("run_tool",))[0]
- result = run_tool.invoke(
- {
- "tool_id": "fake_generate",
- "params": {
- "references": [
- f"https://example.test/{index}.png"
- for index in range(8)
- ]
- },
- }
- )
- self.assertTrue(result["success"])
- def test_tool_operation_journal_replays_and_rejects_drift(self) -> None:
- class CountingClient(FakeRemoteToolClient):
- def __init__(self) -> None:
- self.calls = 0
- def invoke(
- self,
- tool_id: str,
- payload: dict[str, Any],
- ) -> dict[str, Any]:
- self.calls += 1
- return super().invoke(tool_id, payload)
- client = CountingClient()
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- kwargs = {
- "output_dir": root / "outputs",
- "run_dir": root,
- "executor_run_id": "executor-Task1-v1",
- "remote_client": client,
- "allowed_remote_tool_ids": {"fake_generate"},
- "inspected_remote_tool_ids": {"fake_generate"},
- }
- first = create_default_tool_registry(**kwargs)
- first_result = first.resolve(("run_tool",))[0].invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "first"},
- }
- )
- replay = create_default_tool_registry(**kwargs)
- replay_result = replay.resolve(("run_tool",))[0].invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "first"},
- }
- )
- self.assertTrue(first_result["success"])
- self.assertTrue(replay_result["_operation_replayed"])
- self.assertEqual(client.calls, 1)
- drifted = create_default_tool_registry(**kwargs)
- with self.assertRaises(VersionConflictError):
- drifted.resolve(("run_tool",))[0].invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "changed"},
- }
- )
- def test_unknown_tool_outcome_is_not_resent_after_restart(self) -> None:
- session = FakeSession()
- client = RemoteToolClient(
- base_url="https://tools.test/v1/tools",
- session=session,
- )
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- kwargs = {
- "output_dir": root / "outputs",
- "run_dir": root,
- "executor_run_id": "executor-Task1-v1",
- "remote_client": client,
- "allowed_remote_tool_ids": {"fake_generate"},
- "inspected_remote_tool_ids": {"fake_generate"},
- }
- first = create_default_tool_registry(**kwargs)
- with self.assertRaises(OperationOutcomeUnknownError):
- first.resolve(("run_tool",))[0].invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "test"},
- }
- )
- requests_after_first = len(session.requests)
- restarted = create_default_tool_registry(**kwargs)
- with self.assertRaises(OperationOutcomeUnknownError):
- restarted.resolve(("run_tool",))[0].invoke(
- {
- "tool_id": "fake_generate",
- "params": {"prompt": "test"},
- }
- )
- self.assertEqual(len(session.requests), requests_after_first)
- def test_publish_file_local_fallback_and_configured_oss(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- path = Path(temp_dir) / "artifact.txt"
- path.write_text("artifact", encoding="utf-8")
- with patch.dict(
- "os.environ",
- {
- "ALIYUN_OSS_ACCESS_KEY_ID": "",
- "ALIYUN_OSS_ACCESS_KEY_SECRET": "",
- },
- clear=False,
- ):
- self.assertEqual(
- publish_file(path, category="test"),
- str(path.resolve()),
- )
- uploaded: list[tuple[str, str]] = []
- class FakeBucket:
- def put_object_from_file(self, key: str, source: str) -> None:
- uploaded.append((key, source))
- fake_oss = SimpleNamespace(
- Auth=lambda *_: object(),
- Bucket=lambda *_: FakeBucket(),
- )
- with (
- patch.dict(
- "os.environ",
- {
- "ALIYUN_OSS_ACCESS_KEY_ID": "id",
- "ALIYUN_OSS_ACCESS_KEY_SECRET": "secret",
- "ALIYUN_OSS_CDN_BASE_URL": "https://cdn.test",
- },
- clear=False,
- ),
- patch.dict(sys.modules, {"oss2": fake_oss}),
- ):
- url = publish_file(path, category="test")
- self.assertTrue(url.startswith("https://cdn.test/"))
- self.assertEqual(uploaded[0][1], str(path))
- def test_oss_exception_is_treated_as_unknown_side_effect(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- path = Path(temp_dir) / "artifact.txt"
- path.write_text("artifact", encoding="utf-8")
- class FailingBucket:
- def put_object_from_file(self, *_: str) -> None:
- raise TimeoutError("response lost")
- fake_oss = SimpleNamespace(
- Auth=lambda *_: object(),
- Bucket=lambda *_: FailingBucket(),
- )
- with (
- patch.dict(
- "os.environ",
- {
- "ALIYUN_OSS_ACCESS_KEY_ID": "id",
- "ALIYUN_OSS_ACCESS_KEY_SECRET": "secret",
- },
- clear=False,
- ),
- patch.dict(sys.modules, {"oss2": fake_oss}),
- self.assertRaises(UncertainSideEffectError),
- ):
- publish_file(path, category="test")
- class DeterministicImageToolTest(unittest.TestCase):
- def test_image_operations_write_real_files(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- first = root / "first.png"
- second = root / "second.png"
- Image.new("RGB", (100, 80), "red").save(first)
- Image.new("RGB", (100, 80), "blue").save(second)
- crop = crop_image(
- str(first),
- [0, 0, 0.5, 0.5],
- output_dir=root,
- scale=2,
- )
- collage = grid_collage(
- [str(first), str(second)],
- output_dir=root,
- rows=1,
- cols=2,
- )
- text = overlay_text(
- str(first),
- [{"content": "测试", "box": [0.1, 0.1, 0.9, 0.9]}],
- output_dir=root,
- )
- for result in (crop, collage, text):
- self.assertTrue(result["success"])
- self.assertTrue(Path(result["local_paths"][0]).is_file())
- class DeterministicVideoToolTest(unittest.TestCase):
- def _make_video(self, path: Path, color: str) -> None:
- subprocess.run(
- [
- "ffmpeg",
- "-y",
- "-f",
- "lavfi",
- "-i",
- f"color=c={color}:s=64x64:d=1:r=12",
- "-c:v",
- "libx264",
- "-pix_fmt",
- "yuv420p",
- str(path),
- ],
- check=True,
- capture_output=True,
- )
- def _make_audio(self, path: Path) -> None:
- subprocess.run(
- [
- "ffmpeg",
- "-y",
- "-f",
- "lavfi",
- "-i",
- "sine=frequency=440:duration=2",
- str(path),
- ],
- check=True,
- capture_output=True,
- )
- def test_video_operations_and_frames_write_real_files(self) -> None:
- with tempfile.TemporaryDirectory() as temp_dir:
- root = Path(temp_dir)
- first = root / "first.mp4"
- second = root / "second.mp4"
- audio = root / "tone.wav"
- self._make_video(first, "red")
- self._make_video(second, "blue")
- self._make_audio(audio)
- trimmed_audio = trim_audio(
- str(audio),
- output_dir=root,
- start_sec=0.25,
- duration_sec=1.25,
- )
- fitted_audio = fit_audio_duration(
- str(audio),
- target_duration_sec=1.75,
- output_dir=root,
- )
- mixed_audio = mix_audio_tracks(
- [
- {"audio": str(audio), "start_sec": 0, "volume": 1},
- {
- "audio": str(audio),
- "start_sec": 0.5,
- "volume": 0.158,
- },
- ],
- duration_sec=2,
- output_dir=root,
- )
- trimmed = trim_video(
- str(first),
- output_dir=root,
- duration_sec=0.5,
- )
- concatenated = concat_videos(
- [str(first), str(second)],
- output_dir=root,
- )
- muxed = mux_audio(
- str(first),
- str(audio),
- output_dir=root,
- )
- frames = extract_frames(
- str(first),
- output_dir=root,
- num_frames=3,
- )
- silent_info = probe_media(str(first), output_dir=root)
- muxed_info = probe_media(
- muxed["local_paths"][0],
- output_dir=root,
- )
- for result in (
- trimmed_audio,
- fitted_audio,
- mixed_audio,
- trimmed,
- concatenated,
- muxed,
- ):
- self.assertTrue(result["success"])
- self.assertTrue(Path(result["local_paths"][0]).is_file())
- self.assertAlmostEqual(trimmed_audio["duration_sec"], 1.25, places=2)
- self.assertAlmostEqual(fitted_audio["duration_sec"], 1.75, places=2)
- self.assertAlmostEqual(mixed_audio["duration_sec"], 2, places=2)
- self.assertFalse(silent_info["has_audio"])
- self.assertTrue(muxed_info["has_audio"])
- self.assertEqual(muxed_info["audio_codec"], "aac")
- self.assertEqual(len(frames["frames"]), 3)
- self.assertTrue(
- all(
- Path(frame["local_path"]).is_file()
- for frame in frames["frames"]
- )
- )
- if __name__ == "__main__":
- unittest.main()
|