From 1bad97c41acd0136bae9ef9a5c8007d518e80449 Mon Sep 17 00:00:00 2001
From: Jaap Roes <jroes@leukeleu.nl>
Date: Tue, 15 Nov 2022 16:36:46 +0100
Subject: [PATCH] Use configured success_url

Restore the ability to override the success_url in urls.py
---
 tests/test_views_setup.py | 19 +++++++++++++++++++
 tests/urls.py             |  8 ++++++--
 two_factor/views/core.py  |  2 +-
 3 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/tests/test_views_setup.py b/tests/test_views_setup.py
index 06b9bb3f5..497fe371b 100644
--- a/tests/test_views_setup.py
+++ b/tests/test_views_setup.py
@@ -105,6 +105,25 @@ def test_setup_generator_with_multi_method(self):
         self.assertRedirects(response, reverse('two_factor:setup_complete'))
         self.assertEqual(1, self.user.totpdevice_set.count())
 
+    @method_registry(['generator'])
+    def test_setup_custom_success_url(self):
+        custom_setup = reverse('setup-backup_tokens-redirect')
+        custom_redirect = reverse('two_factor:backup_tokens')
+        response = self.client.post(
+            custom_setup,
+            data={
+                'setup_view-current_step': 'method',
+                'method-method': 'generator'
+            }
+        )
+        key = response.context_data['keys'].get('generator')
+        data = {
+            'setup_view-current_step': 'generator',
+            'generator-token': totp(unhexlify(key.encode()))
+        }
+        response = self.client.post(custom_setup, data=data)
+        self.assertRedirects(response, custom_redirect)
+
     def _post(self, data):
         return self.client.post(reverse('two_factor:setup'), data=data)
 
diff --git a/tests/urls.py b/tests/urls.py
index 4e6b9cb56..0b02429fb 100644
--- a/tests/urls.py
+++ b/tests/urls.py
@@ -3,7 +3,7 @@
 
 from two_factor.gateways.twilio.urls import urlpatterns as tf_twilio_urls
 from two_factor.urls import urlpatterns as tf_urls
-from two_factor.views import LoginView
+from two_factor.views import LoginView, SetupView
 
 from .views import SecureView, plain_view
 
@@ -37,7 +37,11 @@
         ),
         name='custom-redirect-authenticated-user-login',
     ),
-
+    path(
+        'account/setup-backup-tokens-redirect/',
+        SetupView.as_view(success_url='two_factor:backup_tokens'),
+        name='setup-backup_tokens-redirect'
+    ),
     path(
         'plain/',
         plain_view,
diff --git a/two_factor/views/core.py b/two_factor/views/core.py
index ebf539c17..69f1aa401 100644
--- a/two_factor/views/core.py
+++ b/two_factor/views/core.py
@@ -446,7 +446,7 @@ class SetupView(RedirectURLMixin, IdempotentSessionWizardView):
     # https://github.com/django/django/blob/58df8aa40fe88f753ba79e091a52f236246260b3/django/contrib/auth/views.py#L63
     def get_success_url(self):
         url = self.get_redirect_url()
-        return url or reverse('two_factor:setup_complete')
+        return url or reverse(self.success_url)
 
     # Copied from django.contrib.auth.views.LoginView (Branch: stable/1.11.x)
     # https://github.com/django/django/blob/58df8aa40fe88f753ba79e091a52f236246260b3/django/contrib/auth/views.py#L67