perf: 添加测试多个账号的任务

apps/assets/api/accounts.py

@@ -3,6 +3,7 @@ from django.conf import settings
 from rest_framework.decorators import action
 from django_filters import rest_framework as filters
 from rest_framework.response import Response
+from rest_framework.generics import CreateAPIView
 
 from orgs.mixins.api import OrgBulkModelViewSet
 from common.permissions import IsOrgAdmin, IsOrgAdminOrAppUser, NeedMFAVerify
@@ -11,7 +12,7 @@ from ..tasks.account_connectivity import test_accounts_connectivity_manual
 from ..models import AuthBook
 from .. import serializers
 
-__all__ = ['AccountViewSet', 'AccountSecretsViewSet']
+__all__ = ['AccountViewSet', 'AccountSecretsViewSet', 'AccountTaskCreateAPI']
 
 
 class AccountFilterSet(BaseFilterSet):
@@ -38,8 +39,6 @@ class AccountFilterSet(BaseFilterSet):
             'asset', 'systemuser', 'id',
         ]
 
-from rest_framework.filters import  SearchFilter
-
 
 class AccountViewSet(OrgBulkModelViewSet):
     model = AuthBook
@@ -79,3 +78,29 @@ class AccountSecretsViewSet(AccountViewSet):
         if not settings.SECURITY_VIEW_AUTH_NEED_MFA:
             self.permission_classes = [IsOrgAdminOrAppUser]
         return super().get_permissions()
+
+
+class AccountTaskCreateAPI(CreateAPIView):
+    permission_classes = (IsOrgAdminOrAppUser,)
+    serializer_class = serializers.AccountTaskSerializer
+    filterset_fields = AccountViewSet.filterset_fields
+    search_fields = AccountViewSet.search_fields
+    filterset_class = AccountViewSet.filterset_class
+
+    def get_accounts(self):
+        queryset = AuthBook.objects.all()
+        queryset = self.filter_queryset(queryset)
+        return queryset
+
+    def perform_create(self, serializer):
+        accounts = self.get_accounts()
+        task = test_accounts_connectivity_manual.delay(accounts)
+        data = getattr(serializer, '_data', {})
+        data["task"] = task.id
+        setattr(serializer, '_data', data)
+        return task
+
+    def get_exception_handler(self):
+        def handler(e, context):
+            return Response({"error": str(e)}, status=400)
+        return handler

apps/assets/serializers/account.py

@@ -40,3 +40,11 @@ class AccountSecretSerializer(AccountSerializer):
             'private_key': {'write_only': False},
             'public_key': {'write_only': False},
         }
+
+
+class AccountTaskSerializer(serializers.Serializer):
+    ACTION_CHOICES = (
+        ('test', 'test'),
+    )
+    action = serializers.ChoiceField(choices=ACTION_CHOICES, write_only=True)
+    task = serializers.CharField(read_only=True)

apps/assets/tasks/account_connectivity.py

@@ -105,3 +105,4 @@ def test_accounts_connectivity_manual(accounts):
     for account in accounts:
         task_name = _("Test account connectivity: {}").format(account)
         test_account_connectivity_util(account, task_name)
+        print(".\n")

apps/assets/urls/api_urls.py

@@ -45,6 +45,8 @@ urlpatterns = [
     path('system-users/<uuid:pk>/tasks/', api.SystemUserTaskApi.as_view(), name='system-user-task-create'),
     path('system-users/<uuid:pk>/cmd-filter-rules/', api.SystemUserCommandFilterRuleListApi.as_view(), name='system-user-cmd-filter-rule-list'),
 
+    path('accounts/tasks/', api.AccountTaskCreateAPI.as_view(), name='account-task-create'),
+
     path('nodes/tree/', api.NodeListAsTreeApi.as_view(), name='node-tree'),
     path('nodes/children/tree/', api.NodeChildrenAsTreeApi.as_view(), name='node-children-tree'),
     path('nodes/<uuid:pk>/children/', api.NodeChildrenApi.as_view(), name='node-children'),