zcbot/tests/test_seedance.py

import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
from uuid import uuid4

import tools.seedance as seedance_module
from tools.seedance import SeedanceTool


class _FakeArkClient:
    submitted_body = None

    def __init__(self, *_args, **_kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return False

    def post_json(self, _endpoint, body, **_kwargs):
        type(self).submitted_body = body
        return {"id": "cgt-test"}

    def get_json(self, _url, **_kwargs):
        return {
            "status": "succeeded",
            "content": {"video_url": "https://example.com/result.mp4"},
        }

    def download(self, _url, dest, **_kwargs):
        Path(dest).write_bytes(b"video")


class SeedanceToolTest(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.root = Path(self._tmp.name) / "user"
        self.root.mkdir()
        _FakeArkClient.submitted_body = None

    def tearDown(self):
        self._tmp.cleanup()

    def make_tool(self) -> SeedanceTool:
        return SeedanceTool(
            ark_cfg=object(),
            video_variant_cfg={
                "model_id": "doubao-seedance-test",
                "default_resolution": "720p",
                "default_ratio": "16:9",
                "default_duration": 5,
                "default_watermark": False,
                "default_generate_audio": False,
                "price_cny_per_mtoken_text2video": 37.0,
                "price_cny_per_mtoken_video2video": 22.0,
                "fps": 24,
                "poll_interval_s": 0,
            },
            variant_key="seedance_test",
            working_dir=self.root,
            user_root=self.root,
            task_id=uuid4(),
            user_id=uuid4(),
        )

    def test_image_to_video_builds_first_frame_and_uses_i2v_price(self):
        (self.root / "source.png").write_bytes(b"\x89PNG\r\n\x1a\n")
        usage = {}

        with (
            patch.object(seedance_module, "ArkClient", _FakeArkClient),
            patch.object(
                seedance_module,
                "record_usage_safe",
                side_effect=lambda _name, _fn, **kw: usage.update(kw),
            ),
        ):
            result = self.make_tool().execute(prompt="让水流动起来", image="source.png")

        content = _FakeArkClient.submitted_body["content"]
        self.assertEqual(content[0], {"type": "text", "text": "让水流动起来"})
        self.assertEqual(content[1]["type"], "image_url")
        self.assertEqual(content[1]["role"], "first_frame")
        self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/png;base64,"))
        self.assertIs(usage["has_video_input"], True)
        self.assertEqual(usage["price_cny_per_mtoken"], 22.0)
        self.assertIn("mode=image_to_video", result)
        self.assertIn("image=source.png", result)

    def test_text_to_video_remains_backward_compatible(self):
        usage = {}

        with (
            patch.object(seedance_module, "ArkClient", _FakeArkClient),
            patch.object(
                seedance_module,
                "record_usage_safe",
                side_effect=lambda _name, _fn, **kw: usage.update(kw),
            ),
        ):
            result = self.make_tool().execute(prompt="水泥浆缓慢流动")

        self.assertEqual(
            _FakeArkClient.submitted_body["content"],
            [{"type": "text", "text": "水泥浆缓慢流动"}],
        )
        self.assertIs(usage["has_video_input"], False)
        self.assertEqual(usage["price_cny_per_mtoken"], 37.0)
        self.assertIn("mode=text_to_video", result)

    def test_multi_image_references_use_reference_role(self):
        for name in ("product-front.png", "product-side.png"):
            (self.root / name).write_bytes(b"\x89PNG\r\n\x1a\n")
        usage = {}

        with (
            patch.object(seedance_module, "ArkClient", _FakeArkClient),
            patch.object(
                seedance_module,
                "record_usage_safe",
                side_effect=lambda _name, _fn, **kw: usage.update(kw),
            ),
        ):
            result = self.make_tool().execute(
                prompt="参考图片1和图片2生成产品环绕展示",
                reference_images=["product-front.png", "product-side.png"],
            )

        content = _FakeArkClient.submitted_body["content"]
        self.assertEqual([part.get("role") for part in content[1:]], [
            "reference_image",
            "reference_image",
        ])
        self.assertEqual(
            usage["extra_units"]["reference_images"],
            ["product-front.png", "product-side.png"],
        )
        self.assertIn("mode=image_to_video", result)
        self.assertIn("images=2", result)

    def test_rejects_more_than_nine_images(self):
        result = self.make_tool().execute(
            prompt="多图参考",
            reference_images=[f"image-{i}.png" for i in range(10)],
        )

        self.assertEqual(
            result,
            "[Error] Seedance 最多支持 9 张输入图片(image + reference_images 合计)。",
        )

    def test_rejects_image_outside_user_root(self):
        outside = self.root.parent / "outside.png"
        outside.write_bytes(b"\x89PNG\r\n\x1a\n")

        result = self.make_tool().execute(prompt="动起来", image=str(outside))

        self.assertTrue(result.startswith("[Error] 图片找不到或越界:"))


if __name__ == "__main__":
    unittest.main()