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

from __future__ import annotations

import hashlib
import json
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

from langchain_core.messages import AIMessage, ToolMessage

from production_build_agents.agents.tool_trace import build_tool_call_records
from production_build_agents.run.operation_journal import (
    OperationOutcomeUnknownError,
)
from production_build_agents.tools import publishing
from production_build_agents.tools.publishing import (
    PublishingToolError,
    publish_media_reference,
    require_media_reference_publishing,
)
from production_build_agents.tools.registry import create_default_tool_registry


_OSS_ENV = {
    "ALIYUN_OSS_ACCESS_KEY_ID": "id",
    "ALIYUN_OSS_ACCESS_KEY_SECRET": "secret",
    "ALIYUN_OSS_ENDPOINT": "https://oss.test",
    "ALIYUN_OSS_BUCKET": "bucket",
    "ALIYUN_OSS_CDN_BASE_URL": "https://cdn.test",
}


class _Download:
    def __init__(self, content: bytes) -> None:
        self.content = content

    def raise_for_status(self) -> None:
        return None

    def iter_content(self, _: int) -> list[bytes]:
        return [self.content]


class SegmentMediaPublishingTest(unittest.TestCase):
    def test_verified_mapping_is_extracted_into_tool_trace(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            local_path = str((Path(temp_dir) / "anchor.png").resolve())
            Path(local_path).write_bytes(b"anchor")
            records = build_tool_call_records(
                [
                    AIMessage(
                        content="",
                        tool_calls=[
                            {
                                "id": "publish-1",
                                "name": "publish_media_reference",
                                "args": {"local_path": local_path},
                            }
                        ],
                    ),
                    ToolMessage(
                        tool_call_id="publish-1",
                        content=json.dumps(
                            {
                                "success": True,
                                "local_path": local_path,
                                "published_url": "https://cdn.test/anchor.png",
                                "content_sha256": "a" * 64,
                            }
                        ),
                    ),
                ]
            )

        self.assertEqual(len(records[0].artifact_mappings), 1)
        self.assertEqual(
            records[0].artifact_mappings[0].source,
            "https://cdn.test/anchor.png",
        )
        self.assertEqual(
            records[0].artifact_mappings[0].local_path,
            local_path,
        )

    def test_content_hash_produces_stable_key_and_verified_mapping(self) -> None:
        uploads: list[tuple[str, str]] = []

        class Bucket:
            def put_object_from_file(self, key: str, source: str) -> None:
                uploads.append((key, source))

        fake_oss = SimpleNamespace(
            Auth=lambda *_: object(),
            Bucket=lambda *_: Bucket(),
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "anchor.png"
            path.write_bytes(b"verified-anchor")
            digest = hashlib.sha256(path.read_bytes()).hexdigest()
            with (
                patch.dict("os.environ", _OSS_ENV, clear=False),
                patch.dict(sys.modules, {"oss2": fake_oss}),
                patch.object(
                    publishing._SESSION,
                    "get",
                    return_value=_Download(path.read_bytes()),
                ),
            ):
                first = publish_media_reference(
                    path,
                    category="segment anchor",
                )
                second = publish_media_reference(
                    path,
                    category="segment anchor",
                )

        self.assertEqual(first, second)
        self.assertEqual(first["local_path"], str(path.resolve()))
        self.assertEqual(first["content_sha256"], digest)
        self.assertEqual(uploads[0][0], uploads[1][0])
        self.assertIn(f"/segment-anchor/{digest}.png", first["published_url"])

    def test_missing_oss_configuration_has_no_local_fallback(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            path = root / "anchor.png"
            path.write_bytes(b"anchor")
            with patch.dict(
                "os.environ",
                {name: "" for name in _OSS_ENV},
                clear=False,
            ):
                registry = create_default_tool_registry(output_dir=root)
                result = registry.resolve(
                    ("publish_media_reference",)
                )[0].invoke({"local_path": str(path)})

        self.assertFalse(result["success"])
        self.assertIn("ALIYUN_OSS_", result["error"])
        self.assertNotIn("published_url", result)
        self.assertNotEqual(result.get("published_url"), str(path.resolve()))

    def test_publishing_preflight_requires_all_formal_oss_settings(self) -> None:
        with (
            patch.dict(
                "os.environ",
                {name: "" for name in _OSS_ENV},
                clear=False,
            ),
            self.assertRaisesRegex(
                PublishingToolError,
                "ALIYUN_OSS_ACCESS_KEY_ID",
            ),
        ):
            require_media_reference_publishing()

    def test_downloaded_bytes_must_match_uploaded_file(self) -> None:
        class Bucket:
            def put_object_from_file(self, *_: str) -> None:
                return None

        fake_oss = SimpleNamespace(
            Auth=lambda *_: object(),
            Bucket=lambda *_: Bucket(),
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            path = Path(temp_dir) / "anchor.png"
            path.write_bytes(b"expected")
            with (
                patch.dict("os.environ", _OSS_ENV, clear=False),
                patch.dict(sys.modules, {"oss2": fake_oss}),
                patch.object(
                    publishing._SESSION,
                    "get",
                    return_value=_Download(b"replaced"),
                ),
                self.assertRaisesRegex(PublishingToolError, "内容与本地"),
            ):
                publish_media_reference(path, category="anchor")

    def test_unknown_upload_outcome_is_not_resent_after_restart(self) -> None:
        attempts = 0

        class Bucket:
            def put_object_from_file(self, *_: str) -> None:
                nonlocal attempts
                attempts += 1
                raise TimeoutError("response lost")

        fake_oss = SimpleNamespace(
            Auth=lambda *_: object(),
            Bucket=lambda *_: Bucket(),
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            path = root / "anchor.png"
            path.write_bytes(b"anchor")
            kwargs = {
                "output_dir": root / "outputs",
                "run_dir": root,
                "executor_run_id": "segment-publish-v1",
            }
            with (
                patch.dict("os.environ", _OSS_ENV, clear=False),
                patch.dict(sys.modules, {"oss2": fake_oss}),
            ):
                first = create_default_tool_registry(**kwargs)
                with self.assertRaises(OperationOutcomeUnknownError):
                    first.resolve(("publish_media_reference",))[0].invoke(
                        {"local_path": str(path), "category": "anchor"}
                    )
                restarted = create_default_tool_registry(**kwargs)
                with self.assertRaises(OperationOutcomeUnknownError):
                    restarted.resolve(
                        ("publish_media_reference",)
                    )[0].invoke(
                        {"local_path": str(path), "category": "anchor"}
                    )

        self.assertEqual(attempts, 1)


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