xiaoxia-saas/tests/unit/test_asset_usage_count_1743.py

"""#1743 素材使用次数回写测试。

计数口径修复：
- mark_asset_used_for_generation 支持 times 参数，按成片实际片段引用次数累加
- _count_plan_clip_asset_usage 从最终成片 plan 的 clips 统计 {asset_id: 片段引用次数}
- _record_rendered_asset_usage 回写 metadata：plan clips 为准、空 plan 兜底任务 asset_ids、
  未被 plan 选用的素材不计数、单素材失败不影响其他素材
"""

from __future__ import annotations

import sys
from pathlib import Path
from unittest.mock import MagicMock

REPO_ROOT = Path(__file__).resolve().parents[2]
for sub in ("apps/worker", "apps/api", "packages", ""):
    p = str(REPO_ROOT / sub) if sub else str(REPO_ROOT)
    if p not in sys.path:
        sys.path.insert(0, p)

import worker_app.tasks.generation as gen_mod  # noqa: E402
from worker_app.core.asset_usage import mark_asset_used_for_generation  # noqa: E402


class FakeAsset:
    def __init__(self, aid: str, metadata: dict | None = None):
        self.id = aid
        self.metadata = metadata or {}


# ── mark_asset_used_for_generation ──────────────────────────────────────────


class TestMarkAssetUsed:
    def test_default_times_is_one(self):
        a = FakeAsset("a1", metadata={})
        mark_asset_used_for_generation(a)
        assert a.metadata["generation_use_count"] == 1
        assert "last_used_at" in a.metadata

    def test_times_accumulates_by_clip_count(self):
        """同一素材被 3 个片段引用 → 一次回写 +3。"""
        a = FakeAsset("a1", metadata={"generation_use_count": 2})
        mark_asset_used_for_generation(a, times=3)
        assert a.metadata["generation_use_count"] == 5

    def test_times_zero_or_negative_floored_to_one(self):
        a = FakeAsset("a1", metadata={})
        mark_asset_used_for_generation(a, times=0)
        assert a.metadata["generation_use_count"] == 1

    def test_preserves_existing_metadata(self):
        a = FakeAsset("a1", metadata={"tags": ["travel"], "generation_use_count": 4})
        mark_asset_used_for_generation(a, times=2)
        assert a.metadata["tags"] == ["travel"]
        assert a.metadata["generation_use_count"] == 6


# ── _count_plan_clip_asset_usage ────────────────────────────────────────────


def _mock_session_with_clips(clip_asset_ids: list[str]):
    """构造 mock session：query(EditPlanClipModel).filter().all() 返回片段 asset_id 行。"""
    rows = [(aid,) for aid in clip_asset_ids]
    session = MagicMock()
    session.query.return_value.filter.return_value.all.return_value = rows
    return session


class TestCountPlanClipUsage:
    def test_counts_asset_occurrences_across_clips(self):
        """plan clips 中 asset-x 出现 2 次、asset-y 1 次 → {x:2, y:1}。"""
        session = _mock_session_with_clips(["asset-x", "asset-y", "asset-x", ""])
        with _patch_clip_model():
            counts = gen_mod._count_plan_clip_asset_usage(session, "plan-1")
        assert counts == {"asset-x": 2, "asset-y": 1}

    def test_empty_clips_returns_empty(self):
        session = _mock_session_with_clips([])
        with _patch_clip_model():
            counts = gen_mod._count_plan_clip_asset_usage(session, "plan-9")
        assert counts == {}

    def test_blank_asset_ids_skipped(self):
        session = _mock_session_with_clips(["", None, "asset-z"])
        with _patch_clip_model():
            counts = gen_mod._count_plan_clip_asset_usage(session, "plan-1")
        assert counts == {"asset-z": 1}


class _patch_clip_model:
    """patch generation 模块内 EditPlanClipModel 的导入路径（函数内 import）。"""

    def __enter__(self):
        import unittest.mock as mock

        self._patches = [
            mock.patch("packages.adapters.sqlalchemy_impl.models.EditPlanClipModel", MagicMock()),
        ]
        for p in self._patches:
            p.start()
        return self

    def __exit__(self, *exc):
        for p in self._patches:
            p.stop()
        return False


# ── _record_rendered_asset_usage ────────────────────────────────────────────


class TestRecordRenderedAssetUsage:
    def _run(self, clip_ids, fallback_ids=None, repo_get_override=None, repo_update_side=None):
        session = _mock_session_with_clips(clip_ids)
        assets: dict[str, FakeAsset] = {}

        def fake_repo_init(sess):
            repo = MagicMock()

            def fake_get(aid):
                if repo_get_override and aid in repo_get_override:
                    return repo_get_override[aid]
                return assets.setdefault(aid, FakeAsset(aid, metadata={}))

            repo.get.side_effect = fake_get
            if repo_update_side:
                repo.update.side_effect = repo_update_side
            else:
                repo.update.side_effect = lambda a: a
            return repo

        import unittest.mock as mock

        with (
            mock.patch(
                "packages.adapters.sqlalchemy_impl.asset_repository.SQLAlchemyAssetRepository",
                side_effect=fake_repo_init,
            ),
            _patch_clip_model(),
        ):
            written = gen_mod._record_rendered_asset_usage(session, "plan-1", "task-1", fallback_asset_ids=fallback_ids)
        return written, assets

    def test_writes_by_plan_clips_not_request_asset_ids(self):
        """核心口径：plan clips 用 x×2 + y×1；请求里的 z（未被 plan 选用）不计数。"""
        written, assets = self._run(
            clip_ids=["asset-x", "asset-x", "asset-y"],
            fallback_ids=["asset-x", "asset-y", "asset-z"],
        )
        assert written == 2
        assert assets["asset-x"].metadata["generation_use_count"] == 2
        assert assets["asset-y"].metadata["generation_use_count"] == 1
        assert "asset-z" not in assets, "未被 plan 选用的素材不应被计数"

    def test_fallback_to_task_asset_ids_when_plan_empty(self):
        """plan 无有效片段（异常数据）→ 兜底任务 asset_ids 每个计 1 次。"""
        written, assets = self._run(clip_ids=[""], fallback_ids=["asset-a", "asset-b"])
        assert written == 2
        assert assets["asset-a"].metadata["generation_use_count"] == 1
        assert assets["asset-b"].metadata["generation_use_count"] == 1

    def test_no_clips_no_fallback_returns_zero(self):
        written, assets = self._run(clip_ids=[], fallback_ids=[])
        assert written == 0
        assert assets == {}

    def test_single_asset_failure_does_not_block_others(self):
        """单素材 update 抛异常不影响其他素材回写。"""

        def update_side(a):
            if a.id == "bad":
                raise RuntimeError("DB boom")
            return a

        written, assets = self._run(
            clip_ids=["good1", "bad", "good2"],
            fallback_ids=None,
            repo_update_side=update_side,
        )
        assert written == 2
        assert assets["good1"].metadata["generation_use_count"] == 1
        assert assets["good2"].metadata["generation_use_count"] == 1

    def test_missing_asset_skipped(self):
        """repo.get 返回 None 的素材跳过，不报错。"""
        written, assets = self._run(
            clip_ids=["ghost", "real"],
            fallback_ids=None,
            repo_get_override={"ghost": None},
        )
        assert written == 1
        assert assets["real"].metadata["generation_use_count"] == 1