yingqing/core/backend/apps/assets/tests.py

from io import BytesIO
from types import SimpleNamespace
from unittest.mock import patch

from django.test import TestCase
from rest_framework.test import APIClient

from apps.accounts.models import Team, TeamMember, User

from .models import Asset, Model


def _mk_team(username, team_name):
    user = User.objects.create_user(username=username, password="x")
    team = Team.objects.create(name=team_name, owner=user)
    TeamMember.objects.create(team=team, user=user, role=TeamMember.Role.OWNER)
    return user, team


class ModelLibraryApiTests(TestCase):
    """模特库:本团队模特 ∪ 官方模板;官方跨团队可见;软删;官方不可删;团队隔离。"""

    def setUp(self):
        self.user, self.teamA = _mk_team("ua", "TeamA")
        self.other, self.teamB = _mk_team("ub", "TeamB")
        self.client = APIClient()
        self.client.force_authenticate(self.user)

        self.mine = Model.objects.create(team=self.teamA, name="我的模特")
        self.official = Model.objects.create(team=self.teamB, name="官方小美", is_official=True)
        self.bs_other = Model.objects.create(team=self.teamB, name="别队私有")  # B 的私有,A 看不到

    def test_list_returns_mine_plus_official(self):
        ids = {m["id"] for m in self.client.get("/api/models/").json()["results"]}
        self.assertIn(str(self.mine.id), ids)
        self.assertIn(str(self.official.id), ids)  # 官方跨团队可见
        self.assertNotIn(str(self.bs_other.id), ids)  # 别队私有不可见

    def test_tab_official_and_mine(self):
        off = {m["id"] for m in self.client.get("/api/models/?tab=official").json()["results"]}
        self.assertEqual(off, {str(self.official.id)})
        mine = {m["id"] for m in self.client.get("/api/models/?tab=mine").json()["results"]}
        self.assertEqual(mine, {str(self.mine.id)})

    def test_official_flag_in_payload(self):
        row = next(m for m in self.client.get("/api/models/").json()["results"] if m["id"] == str(self.official.id))
        self.assertTrue(row["is_official"])

    def test_delete_is_soft(self):
        res = self.client.delete(f"/api/models/{self.mine.id}/")
        self.assertEqual(res.status_code, 204)
        self.mine.refresh_from_db()
        self.assertTrue(self.mine.is_deleted)
        ids = {m["id"] for m in self.client.get("/api/models/").json()["results"]}
        self.assertNotIn(str(self.mine.id), ids)

    def test_official_cannot_be_deleted(self):
        res = self.client.delete(f"/api/models/{self.official.id}/")
        self.assertEqual(res.status_code, 403)
        self.official.refresh_from_db()
        self.assertFalse(self.official.is_deleted)

    def test_cannot_delete_other_team_model(self):
        res = self.client.delete(f"/api/models/{self.bs_other.id}/")
        self.assertIn(res.status_code, (403, 404))  # 别队私有:不在可见 qs → 404,或命中守卫 → 403
        self.bs_other.refresh_from_db()
        self.assertFalse(self.bs_other.is_deleted)

    def test_create_rejects_cross_team_asset(self):
        b_asset = Asset.objects.create(team=self.teamB, name="别队图", asset_type="image", source="upload")
        res = self.client.post("/api/models/", {"name": "盗用", "portrait_asset": str(b_asset.id)}, format="json")
        self.assertEqual(res.status_code, 400)

    def test_upload_real_person_creates_asset_and_model(self):
        stored = SimpleNamespace(object_key="teams/x/models/a.png", bucket="b", content_type="image/png", size_bytes=10)
        with patch("apps.assets.views.TosStorage") as Tos:
            Tos.return_value.upload_fileobj.return_value = stored
            f = BytesIO(b"img")
            f.name = "lady.png"
            res = self.client.post("/api/models/upload/", {"name": "上传妹", "file": f}, format="multipart")
        self.assertEqual(res.status_code, 201)
        body = res.json()
        self.assertEqual(body["source"], "upload")
        model = Model.objects.get(id=body["id"])
        self.assertIsNotNone(model.portrait_asset)
        self.assertEqual(model.portrait_asset.category, Asset.Category.MODEL_PORTRAIT)


class AssetSoftDeleteTests(TestCase):
    """资产库不展示软删资产(is_deleted=True):列表 + tab 计数都排除。"""

    def setUp(self):
        self.user, self.team = _mk_team("usd", "TeamSD")
        self.client = APIClient()
        self.client.force_authenticate(self.user)
        self.live = Asset.objects.create(team=self.team, name="留存图", asset_type="image", source="ai_generated", category=Asset.Category.PERSON)
        self.dead = Asset.objects.create(team=self.team, name="废图", asset_type="image", source="ai_generated", category=Asset.Category.PERSON, is_deleted=True)

    def test_list_excludes_soft_deleted(self):
        ids = {a["id"] for a in self.client.get("/api/assets/?page_size=50").json()["results"]}
        self.assertIn(str(self.live.id), ids)
        self.assertNotIn(str(self.dead.id), ids)

    def test_summary_excludes_soft_deleted(self):
        self.assertEqual(self.client.get("/api/assets/summary/").json()["people"], 1)


class ReviewScopeTests(TestCase):
    """送审范围 = 角色定妆照(person)/ 三视图(tri_view)/ 分镜图(storyboard);图片趴不送审。"""

    def test_review_categories_constant(self):
        self.assertEqual(Asset.REVIEW_CATEGORIES, ("person", "tri_view", "storyboard"))

    def test_poll_team_reviews_only_scans_review_categories(self):
        from apps.assets import review

        _, team = _mk_team("uc", "TeamC")
        keep = {}
        for cat in ("person", "tri_view", "storyboard"):
            a = Asset.objects.create(team=team, name=cat, asset_type="image", source="ai_generated",
                                     category=cat, review_status="processing")
            keep[cat] = str(a.id)
        # 图片趴 / 场景:不应进送审队列
        for cat in ("model_tryon", "model_portrait", "scene"):
            Asset.objects.create(team=team, name=cat, asset_type="image", source="ai_generated",
                                 category=cat, review_status="processing")
        with patch("apps.assets.review.assets_client.is_enabled", return_value=False):
            out = review.poll_team_reviews(team)
        self.assertEqual(set(out.keys()), set(keep.values()))