from __future__ import unicode_literals from django.test import TestCase from django.test.client import RequestFactory from rest_framework import status from rest_framework.response import Response from rest_framework import viewsets from rest_framework.decorators import link, action from rest_framework.routers import SimpleRouter import copy factory = RequestFactory() class BasicViewSet(viewsets.ViewSet): def list(self, request, *args, **kwargs): return Response({'method': 'list'}) @action() def action1(self, request, *args, **kwargs): return Response({'method': 'action1'}) @action() def action2(self, request, *args, **kwargs): return Response({'method': 'action2'}) @link() def link1(self, request, *args, **kwargs): return Response({'method': 'link1'}) @link() def link2(self, request, *args, **kwargs): return Response({'method': 'link2'}) class TestSimpleRouter(TestCase): def setUp(self): self.router = SimpleRouter() def test_link_and_action_decorator(self): routes = self.router.get_routes(BasicViewSet) decorator_routes = routes[2:] # Make sure all these endpoints exist and none have been clobbered for i, endpoint in enumerate(['action1', 'action2', 'link1', 'link2']): route = decorator_routes[i] # check url listing self.assertEqual(route.url, '^{{prefix}}/{{lookup}}/{0}/$'.format(endpoint)) # check method to function mapping if endpoint.startswith('action'): method_map = 'post' else: method_map = 'get' self.assertEqual(route.mapping[method_map], endpoint)