测试开发进阶(三十六)

def list(self, request, *args, **kwargs):
    queryset = self.filter_queryset(self.get_queryset())

    page = self.paginate_queryset(queryset)
    if page is not None:
        serializer = self.get_serializer(page, many=True)
        datas = serializer.data
        datas = get_count_by_project(datas)
        return self.get_paginated_response(datas)

    serializer = self.get_serializer(queryset, many=True)
    datas = serializer.data
    datas = get_count_by_project(datas)
    return Response(datas)

super().list(request, *args, **kwargs)

def list(self, request, *args, **kwargs):
    response = super().list(request, *args, **kwargs)
    response.data['results'] = get_count_by_project(response.data['results'])
    return response

def get_serializer_class(self):
    """
    Return the class to use for the serializer.
    Defaults to using `self.serializer_class`.

    You may want to override this if you need to provide different
    serializations depending on the incoming request.

    (Eg. admins get full serialization, others get basic serialization)
    """
    assert self.serializer_class is not None, (
        "'%s' should either include a `serializer_class` attribute, "
        "or override the `get_serializer_class()` method."
        % self.__class__.__name__
    )

    return self.serializer_class

def get_serializer_class(self):
    if self.action == 'names':
        return serializers.ProjectNameSerializer
    else:
        return self.serializer_class

from datetime import datetime
from rest_framework import serializers
from .models import Reports


class ReportsSerializer(serializers.ModelSerializer):
    """
    报告序列化器
    """

    class Meta:
        model = Reports
        exclude = ('update_time', 'is_delete')
        extra_kwargs = {
            'html': {
                'write_only': True
            },
            'create_time': {
                'read_only': True
            }
        }

    def create(self, validated_data):
        report_name = validated_data['name']
        validated_data['name'] = f"{report_name}_{datetime.strftime(datetime.now(), '%Y%m%d%H%M%S')}"
        report = Reports.objects.create(**validated_data)
        return report

name = models.CharField('报告名称', max_length=200, unique=True, help_text='报告名称')

import re
import os
from datetime import datetime

from django.conf import settings
from django.http import StreamingHttpResponse
from rest_framework.viewsets import ModelViewSet
from rest_framework import permissions
from rest_framework.decorators import action

from reports.utils import format_output, get_file_contents
from .models import Reports
from .serializers import ReportsSerializer


class ReportsViewSet(ModelViewSet):
    """
    list:
    返回测试报告（多个）列表数据

    create:
    创建测试报告

    retrieve:
    返回测试报告（单个）详情数据

    update:
    更新（全）测试报告

    partial_update:
    更新（部分）测试报告

    destroy:
    删除测试报告

    """
    queryset = Reports.objects.filter(is_delete=False)
    serializer_class = ReportsSerializer
    permission_classes = (permissions.IsAuthenticated,)
    ordering_fields = ('id', 'name')

    def perform_destroy(self, instance):
        instance.is_delete = True
        instance.save()

    def list(self, request, *args, **kwargs):
        response = super().list(request, *args, **kwargs)
        response.data['results'] = format_output(response.data['results'])
        return response

    @action(detail=True)
    def download(self, request, pk=None):
        instance = self.get_object()
        html = instance.html
        name = instance.name
        mtch = re.match(r'(.*_)\d+', name)
        if mtch:
            mtch = mtch.group(1) + datetime.strftime(datetime.now(), '%Y%m%d%H%M%S') + '.html'
        report_dir = os.path.join(settings.BASE_DIR, 'reports')
        report_path = os.path.join(report_dir, mtch)
        with open(report_path, 'w') as f:
            f.write(html)
        response = StreamingHttpResponse(get_file_contents(report_path))
        response['Content-Type'] = "application/octet-stream"
        response['Content-Disposition'] = "attachment; filename*=UTF-8''{}".format(name)
        return response

response = StreamingHttpResponse(get_file_contents(report_path))

def get_file_contents(filename, chunk_size=512):
    with open(filename,encoding='utf8') as f:
        while True:
            c = f.read(chunk_size)
            if c:
                yield c
            else:
                break