comfyui_o1key/tests/test_seedance_asset_service.py

"""Offline tests for shared Seedance asset creation and ID reuse."""

from __future__ import annotations

import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, Mock


PACKAGE_DIR = Path(__file__).resolve().parents[1]
COMFY_ROOT = PACKAGE_DIR.parent.parent
sys.path.insert(0, str(COMFY_ROOT))
sys.path.insert(0, str(PACKAGE_DIR.parent))

package = types.ModuleType("comfyui_o1key")
package.__path__ = [str(PACKAGE_DIR)]
sys.modules["comfyui_o1key"] = package
for child in ("utils", "clients"):
    module = types.ModuleType(f"comfyui_o1key.{child}")
    module.__path__ = [str(PACKAGE_DIR / child)]
    sys.modules[f"comfyui_o1key.{child}"] = module

from comfyui_o1key.utils.seedance_assets import (  # noqa: E402
    SeedanceAssetService,
    seedance_asset_fingerprint,
    seedance_asset_request_type,
)


class SeedanceAssetServiceTests(unittest.IsolatedAsyncioTestCase):
    def test_video_routes_map_to_the_create_node_request_types(self):
        self.assertEqual(seedance_asset_request_type("overseas_hc"), "hc")
        self.assertEqual(seedance_asset_request_type("domestic"), "doubao")
        with self.assertRaisesRegex(ValueError, "素材线路"):
            seedance_asset_request_type("unknown")

    async def test_active_content_cache_skips_upload_and_creation(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            media_path = Path(temp_dir, "reference.bin")
            media_path.write_bytes(b"same-content")
            cache_path = Path(temp_dir, "seedance_asset_cache.json")
            fingerprint = seedance_asset_fingerprint(str(media_path), "doubao", "video")

            first_client = type("FirstClient", (), {})()
            first_client.create_hc_asset_and_wait = AsyncMock(return_value={
                "Id": "asset-safe-id",
                "Status": "Active",
            })
            first_client.get_hc_asset = AsyncMock()
            first_upload = AsyncMock(return_value="https://upload.example.invalid/reference-video")
            first = SeedanceAssetService(client=first_client, cache_path=str(cache_path))
            created = await first.create_from_url(
                name="reference",
                asset_url_factory=first_upload,
                asset_type="video",
                request_type="doubao",
                fingerprint=fingerprint,
            )
            self.assertFalse(created["_reused"])
            first_upload.assert_awaited_once()
            first_client.create_hc_asset_and_wait.assert_awaited_once()

            second_client = type("SecondClient", (), {})()
            second_client.get_hc_asset = AsyncMock(return_value={
                "Id": "asset-safe-id",
                "Status": "Active",
            })
            second_client.create_hc_asset_and_wait = AsyncMock()
            second_upload = AsyncMock(return_value="https://upload.example.invalid/should-not-upload")
            second = SeedanceAssetService(client=second_client, cache_path=str(cache_path))
            reused = await second.create_from_url(
                name="reference",
                asset_url_factory=second_upload,
                asset_type="video",
                request_type="doubao",
                fingerprint=fingerprint,
            )
            self.assertTrue(reused["_reused"])
            self.assertEqual(reused["Id"], "asset-safe-id")
            second_upload.assert_not_awaited()
            second_client.create_hc_asset_and_wait.assert_not_awaited()

            cache_payload = json.loads(cache_path.read_text(encoding="utf-8"))
            serialized = json.dumps(cache_payload, ensure_ascii=False)
            self.assertIn("asset-safe-id", serialized)
            self.assertNotIn("upload.example.invalid", serialized)
            self.assertNotIn(str(media_path), serialized)

            stale_client = type("StaleClient", (), {})()
            stale_client.get_hc_asset = AsyncMock(
                side_effect=RuntimeError("查询 Doubao 素材失败 (404): missing")
            )
            stale_client.create_hc_asset_and_wait = AsyncMock(return_value={
                "Id": "asset-recreated-id",
                "Status": "Active",
            })
            stale_upload = AsyncMock(return_value="https://upload.example.invalid/recreated")
            stale = SeedanceAssetService(client=stale_client, cache_path=str(cache_path))
            recreated = await stale.create_from_url(
                name="reference",
                asset_url_factory=stale_upload,
                asset_type="video",
                request_type="doubao",
                fingerprint=fingerprint,
            )
            self.assertFalse(recreated["_reused"])
            self.assertEqual(recreated["Id"], "asset-recreated-id")
            stale_upload.assert_awaited_once()
            stale_client.create_hc_asset_and_wait.assert_awaited_once()

    async def test_cache_write_failure_does_not_discard_created_asset(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            media_path = Path(temp_dir, "reference.bin")
            media_path.write_bytes(b"content")
            client = type("Client", (), {})()
            client.create_hc_asset_and_wait = AsyncMock(return_value={
                "Id": "asset-created-id",
                "Status": "Active",
            })
            client.get_hc_asset = AsyncMock()
            service = SeedanceAssetService(
                client=client,
                cache_path=str(Path(temp_dir, "cache.json")),
            )
            service.cache.put = Mock(side_effect=OSError("read only"))

            result = await service.create_from_url(
                name="reference",
                asset_url="https://upload.example.invalid/reference",
                asset_type="image",
                request_type="hc",
                fingerprint=seedance_asset_fingerprint(str(media_path), "hc", "image"),
            )

            self.assertEqual(result["Id"], "asset-created-id")


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