video-image-production-build/tests/tools/test_seedance.py

from __future__ import annotations

import unittest
from typing import Any

from production_build_agents.tools.seedance import (
    HuoshanSeedanceClient,
)


class FakeResponse:
    def __init__(
        self,
        payload: dict[str, Any] | None = None,
        *,
        status_code: int = 200,
    ) -> None:
        self.payload = payload or {}
        self.status_code = status_code

    def json(self) -> dict[str, Any]:
        return self.payload


class SeedanceSession:
    def __init__(self) -> None:
        self.submitted_json: dict[str, Any] | None = None
        self.post_count = 0
        self.get_count = 0

    def post(self, _: str, **kwargs: Any) -> FakeResponse:
        self.post_count += 1
        self.submitted_json = kwargs["json"]
        return FakeResponse({"id": "cgt-task-1"})

    def get(self, _: str, **__: Any) -> FakeResponse:
        self.get_count += 1
        if self.get_count == 1:
            return FakeResponse(
                {
                    "id": "cgt-task-1",
                    "status": "running",
                }
            )
        return FakeResponse(
            {
                "id": "cgt-task-1",
                "model": "doubao-seedance-2-0-260128",
                "status": "succeeded",
                "content": {
                    "video_url": "https://result.test/shot.mp4"
                },
                "usage": {"completion_tokens": 42},
            }
        )


def _payload() -> dict[str, Any]:
    return {
        "prompt": "人物面向镜头自然讲述",
        "first_frame_url": "https://input.test/first.jpg",
        "reference_image_urls": ["https://input.test/reference.jpg"],
        "reference_video_urls": ["https://input.test/motion.mp4"],
        "last_frame_url": "https://input.test/last.jpg",
        "duration": 5,
        "ratio": "9:16",
        "resolution": "1080p",
        "generate_audio": False,
        "watermark": False,
        "web_search": False,
    }


class HuoshanSeedanceClientTest(unittest.TestCase):
    def test_maps_stable_tool_contract_and_polls_same_task(self) -> None:
        session = SeedanceSession()
        client = HuoshanSeedanceClient(
            api_key="secret",
            session=session,  # type: ignore[arg-type]
            poll_interval=0,
        )

        result = client.invoke("seedance_generate_video", _payload())

        self.assertTrue(result["success"])
        self.assertEqual(result["data"]["task_id"], "cgt-task-1")
        self.assertEqual(
            result["data"]["result_url"],
            "https://result.test/shot.mp4",
        )
        self.assertEqual(session.post_count, 1)
        self.assertEqual(session.get_count, 2)
        assert session.submitted_json is not None
        self.assertEqual(
            session.submitted_json["model"],
            "doubao-seedance-2-0-260128",
        )
        self.assertEqual(
            [
                item.get("role")
                for item in session.submitted_json["content"][1:]
            ],
            [
                "first_frame",
                "reference_image",
                "reference_video",
                "last_frame",
            ],
        )
        self.assertFalse(session.submitted_json["generate_audio"])
        self.assertFalse(session.submitted_json["watermark"])

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