前言
如果你正在使用 Django Rest Framework 来构建 RESTful API，那么你一定会经常使用到 ViewSet 这个工具类。ViewSet 提供了一些常见操作的映射，比如 list、create、retrieve、update、destroy 等，能够很方便地实现 CRUD 操作。
不过，如果你需要实现一些比较特殊的操作，那么 ViewSet 的默认方法就可能无法满足你的需求。例如，在一个投票系统中，我们需要提供一个对某个选项进行投票的接口，此时 ViewSet 默认的 list、create、retrieve 等方法就已经无法胜任此任务了。
在这种情况下，我们可以使用 Django Rest Framework 提供的 @action 装饰器来自定义方法。
如何使用 @action 定义自定义方法
使用 @action 定义自定义方法很简单，只需要在 ViewSet 中定义一个方法，并在该方法上加上 @action 装饰器，就可以将该方法转换为一个 API 接口。例如，我们可以在一个投票系统的 ViewSet 中定义一个投票接口：
from rest_framework.decorators import action
from rest_framework.response import Response
class PollViewSet(viewsets.ModelViewSet):
queryset = Poll.objects.all()
serializer_class = PollSerializer
@action(methods=['post'], detail=True)
def vote(self, request, pk=None):
poll = self.get_object()
option_id = request.data.get('option', None)
if not option_id:
return Response({'error': 'option is required'})
try:
option = poll.options.get(pk=option_id)
except Option.DoesNotExist:
return Response({'error': 'invalid option'})
option.votes += 1
option.save()
return Response({'success': True})
该投票接口仅支持 POST 请求，它将会根据传入的 option 参数来对某个选项进行投票。这里需要注意的是，需要通过 self.get_object() 方法来获取当前的 Poll 对象，需要通过 request.data 来获取 POST 请求的参数。
在定义完自定义方法之后，需要将该方法添加到路由中，这可以通过 Django Rest Framework 提供的默认路由机制来实现：
from rest_framework import routers
router = routers.DefaultRouter()
router.register(r'polls', PollViewSet)
urlpatterns = [
path('', include(router.urls)),
]
这里的 router.register() 方法将会自动根据 PollViewSet 中定义的方法来生成对应的路由，并将其添加到 urlpatterns 中。
对于 @action 定制方法，Django Rest Framework 会自动生成对应的 URL，也就是说，对于上面的投票例子，URL 是：
/polls/<pk>/vote/
其中 <pk> 对应的是具体的 Poll 对象的 primary key。例如，如果你的 Poll 对象的 primary key 是 3，那么可以这样访问：
/polls/3/vote/
在这个例子中，@action(detail=True) 会自动在 URL 中添加 pk 参数（代表 Poll 对象的 primary key），如果没有设置 detail=True，那么 URL 中就不会出现 pk 参数。由此可以看出，@action(detail=True) 与 @action(detail=False) 的区别在于 URL 中是否需要添加 pk 参数。
至此，我们已经可以使用localhost:8000/polls/<id>/vote来访问这个接口了，如果你在其他定制化的时候不需要传入id参数，那么你可以配置detail=False，这样子就可以使用localhost:8000/polls/vote来访问这个接口了。
测试自定义方法
在实现完自定义方法之后，我们需要对其进行测试，以确保其可以正常工作。这可以通过 Django Rest Framework 提供的 APIClient 来实现。
from django.test import TestCase
from django.urls import reverse
from rest_framework import status
from rest_framework.test import APIClient
from polls.models import Poll, Option
class PollAPITestCase(TestCase):
def setUp(self):
self.client = APIClient()
def test_vote(self):
poll = Poll.objects.create(title='test poll')
option = Option.objects.create(poll=poll, title='option1')
url = reverse('polls-vote', args=[poll.id])
data = {'option': option.id}
response = self.client.post(url, data, format='json')
self.assertEqual(response.status_code, status.HTTP_200_OK)
在这个示例测试中，我们首先使用 Poll.objects.create() 方法创建了一个 Poll 对象和一个 Option 对象，接着使用 Django Rest Framework 提供的 reverse() 函数根据 URL 模式的名称反向生成对应的 URL。最后，我们使用 APIClient 的 post() 方法向该 URL 发送 post 请求，并将选项的主键作为数据传递进去。最终，我们断言响应的 HTTP 状态码是否为 200，以此来判断测试是否通过。
reverse() 函数用来根据 URL 模式的名称生成对应的 URL，它接收一个参数，该参数是 URL 模式的名称，会将该名称反向映射为对应的 URL。例如，reverse('polls-vote', args=[poll.id]) 将会生成 URL /polls/<poll.id>/vote/。如果 URL 中含有不止一个占位符，则可以使用类似于 [poll.id, option.id] 的方式将多个占位符传递到 reverse() 函数中，最终生成具有多个参数的 URL。
总结
在本文中，我们介绍了如何使用 Django Rest Framework 的 @action 装饰器来定义自定义方法，以及如何使用 Django Rest Framework 的 APIClient 对自定义方法进行单元测试。通过本文，你应该已经掌握了 @action 装饰器的用法，以及如何进行单元测试。希望这篇文章能够帮助你更好地理解 Django Rest Framework。
评论区