test: 补全StubGenerationTaskRepository的限流计数方法

2026-07-11 20:36:52 +08:00
parent 4ebbafeb8a
commit 9427b5873d
1 changed files with 10 additions and 0 deletions
@@ -118,6 +118,16 @@ class StubGenerationTaskRepository:
    def count_by_user(self, user_id: str) -> int:
        return len([t for t in self._tasks.values() if t.created_by_user_id == user_id])

+    def count_pending_by_user(self, user_id: str) -> int:
+        from app.models.generation_task import TaskStatus
+        pending_statuses = {TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.QUEUED}
+        return len([t for t in self._tasks.values() if t.created_by_user_id == user_id and t.status in pending_statuses])
+
+    def count_pending_total(self) -> int:
+        from app.models.generation_task import TaskStatus
+        pending_statuses = {TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.QUEUED}
+        return len([t for t in self._tasks.values() if t.status in pending_statuses])
+
    def list_recent_by_user(self, user_id: str, limit: int = 5) -> list[GenerationTask]:
        items = [t for t in self._tasks.values() if t.created_by_user_id == user_id]
        items.sort(key=lambda t: t.created_at, reverse=True)