Merge pull request #192 from PierreF/login-redirect-url-as-url

Fix failure if LOGIN_REDIRECT_URL is an URL

tests/test_views_phone.py

@@ -6,7 +6,7 @@
     import mock
 
 from django.conf import settings
-from django.core.urlresolvers import reverse
+from django.core.urlresolvers import reverse, reverse_lazy
 from django.shortcuts import resolve_url
 from django.test import TestCase
 from django.test.utils import override_settings
@@ -17,6 +17,7 @@
 from two_factor.models import PhoneDevice
 from two_factor.utils import backup_phones
 from two_factor.validators import validate_international_phonenumber
+from two_factor.views.core import PhoneDeleteView, PhoneSetupView
 
 from .utils import UserMixin
 
@@ -79,6 +80,50 @@ def test_number_validation(self, fake):
             response.context_data['wizard']['form'].errors,
             {'number': [six.text_type(validate_international_phonenumber.message)]})
 
+    @mock.patch('formtools.wizard.views.WizardView.get_context_data')
+    def test_success_url_as_url(self, get_context_data):
+        url = '/account/two_factor/'
+        view = PhoneSetupView()
+        view.success_url = url
+
+        def return_kwargs(form, **kwargs):
+            return kwargs
+        get_context_data.side_effect = return_kwargs
+
+        context = view.get_context_data(None)
+        self.assertIn('cancel_url', context)
+        self.assertEqual(url, context['cancel_url'])
+
+    @mock.patch('formtools.wizard.views.WizardView.get_context_data')
+    def test_success_url_as_named_url(self, get_context_data):
+        url_name = 'two_factor:profile'
+        url = reverse(url_name)
+        view = PhoneSetupView()
+        view.success_url = url_name
+
+        def return_kwargs(form, **kwargs):
+            return kwargs
+        get_context_data.side_effect = return_kwargs
+
+        context = view.get_context_data(None)
+        self.assertIn('cancel_url', context)
+        self.assertEqual(url, context['cancel_url'])
+
+    @mock.patch('formtools.wizard.views.WizardView.get_context_data')
+    def test_success_url_as_reverse_lazy(self, get_context_data):
+        url_name = 'two_factor:profile'
+        url = reverse(url_name)
+        view = PhoneSetupView()
+        view.success_url = reverse_lazy(url_name)
+
+        def return_kwargs(form, **kwargs):
+            return kwargs
+        get_context_data.side_effect = return_kwargs
+
+        context = view.get_context_data(None)
+        self.assertIn('cancel_url', context)
+        self.assertEqual(url, context['cancel_url'])
+
 
 class PhoneDeleteTest(UserMixin, TestCase):
     def setUp(self):
@@ -99,6 +144,26 @@ def test_cannot_delete_default(self):
                                             args=[self.default.pk]))
         self.assertContains(response, 'was not found', status_code=404)
 
+    def test_success_url_as_url(self):
+        url = '/account/two_factor/'
+        view = PhoneDeleteView()
+        view.success_url = url
+        self.assertEqual(view.get_success_url(), url)
+
+    def test_success_url_as_named_url(self):
+        url_name = 'two_factor:profile'
+        url = reverse(url_name)
+        view = PhoneDeleteView()
+        view.success_url = url_name
+        self.assertEqual(view.get_success_url(), url)
+
+    def test_success_url_as_reverse_lazy(self):
+        url_name = 'two_factor:profile'
+        url = reverse(url_name)
+        view = PhoneDeleteView()
+        view.success_url = reverse_lazy(url_name)
+        self.assertEqual(view.get_success_url(), url)
+
 
 class PhoneDeviceTest(UserMixin, TestCase):
     def test_verify(self):

two_factor/views/core.py

@@ -11,7 +11,7 @@
 from django.contrib.auth.decorators import login_required
 from django.contrib.auth.forms import AuthenticationForm
 from django.contrib.sites.shortcuts import get_current_site
-from django.core.urlresolvers import reverse, reverse_lazy
+from django.core.urlresolvers import reverse
 from django.forms import Form
 from django.http import Http404, HttpResponse
 from django.shortcuts import redirect, resolve_url
@@ -424,7 +424,7 @@ class PhoneSetupView(IdempotentSessionWizardView):
     numbers can be used for verification.
     """
     template_name = 'two_factor/core/phone_register.html'
-    success_url = reverse_lazy(settings.LOGIN_REDIRECT_URL)
+    success_url = settings.LOGIN_REDIRECT_URL
     form_list = (
         ('setup', PhoneNumberMethodForm),
         ('validation', DeviceValidationForm),
@@ -481,7 +481,7 @@ def get_key(self):
         return self.storage.extra_data[self.key_name]
 
     def get_context_data(self, form, **kwargs):
-        kwargs.setdefault('cancel_url', self.success_url)
+        kwargs.setdefault('cancel_url', resolve_url(self.success_url))
         return super(PhoneSetupView, self).get_context_data(form, **kwargs)
 
 
@@ -491,11 +491,14 @@ class PhoneDeleteView(DeleteView):
     """
     View for removing a phone number used for verification.
     """
-    success_url = reverse_lazy(settings.LOGIN_REDIRECT_URL)
+    success_url = settings.LOGIN_REDIRECT_URL
 
     def get_queryset(self):
         return self.request.user.phonedevice_set.filter(name='backup')
 
+    def get_success_url(self):
+        return resolve_url(self.success_url)
+
 
 @class_view_decorator(never_cache)
 @class_view_decorator(otp_required)