优化脚本就该模型

2026-08-21 15:08:24 +08:00
parent c2d3a786a3
commit 158c66ac25
12 changed files with 273 additions and 11 deletions
@@ -515,6 +515,61 @@ class AdminTaskMonitorTests(TestCase):
        task = self._mk_stuck_reserved("k-reap-perm")
        self.assertEqual(self.nc.post(f"/api/admin/tasks/{task.id}/reap/").status_code, 403)

+    def _mk_inflight_video(self, *, key: str, provider_task_id: str = "ark-poll-1"):
+        return self.AITask.objects.create(
+            team=self.team,
+            model_config=self.mc,
+            task_type=self.AITask.Type.FREE_VIDEO,
+            status=self.AITask.Status.POLLING,
+            estimated_cost="1.0",
+            actual_cost="0",
+            idempotency_key=key,
+            provider_task_id=provider_task_id,
+        )
+
+    def test_poll_inflight_free_video(self):
+        task = self._mk_inflight_video(key="k-poll-fv")
+
+        def fake_finalize(*, task):
+            task.status = self.AITask.Status.SUCCEEDED
+            task.save(update_fields=["status", "updated_at"])
+            return task
+
+        with patch("apps.ai.free_video.finalize_free_video", side_effect=fake_finalize) as mock_fin:
+            r = self.ac.post("/api/admin/tasks/poll/", {}, format="json")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.data["polled"], 1)
+        self.assertEqual(r.data["statuses"][str(task.id)], "succeeded")
+        mock_fin.assert_called_once()
+        task.refresh_from_db()
+        self.assertEqual(task.status, self.AITask.Status.SUCCEEDED)
+
+    def test_poll_skips_terminal_and_missing_provider_id(self):
+        self.AITask.objects.create(
+            team=self.team, model_config=self.mc, task_type=self.AITask.Type.FREE_VIDEO,
+            status=self.AITask.Status.POLLING, estimated_cost="1.0", actual_cost="0",
+            idempotency_key="k-poll-nopid",
+        )
+        with patch("apps.ai.free_video.finalize_free_video") as mock_fin:
+            r = self.ac.post("/api/admin/tasks/poll/", {}, format="json")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.data["polled"], 0)
+        mock_fin.assert_not_called()
+
+    def test_poll_task_ids_filter(self):
+        keep = self._mk_inflight_video(key="k-poll-keep", provider_task_id="ark-keep")
+        skip = self._mk_inflight_video(key="k-poll-skip", provider_task_id="ark-skip")
+        with patch("apps.ai.free_video.finalize_free_video", side_effect=lambda *, task: task) as mock_fin:
+            r = self.ac.post("/api/admin/tasks/poll/", {"task_ids": [str(keep.id)]}, format="json")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(r.data["polled"], 1)
+        self.assertIn(str(keep.id), r.data["statuses"])
+        self.assertNotIn(str(skip.id), r.data["statuses"])
+        mock_fin.assert_called_once()
+
+    def test_poll_requires_admin(self):
+        self.assertEqual(self.nc.post("/api/admin/tasks/poll/", {}, format="json").status_code, 403)
+

 class AdminBillingTests(TestCase):
    """Phase 7:计费审计(流水浏览/手动调额)+ 4 层额度策略(CRUD + 拦截生效)+ 权限。"""