From 349a3a67d55d27db8eb065c1c7fdb79a77b3f8fc Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Thu, 17 Dec 2020 19:31:16 +0100
Subject: [PATCH] flows: use to_stage_response in _flow_done()

---
 authentik/flows/tests/test_views.py | 14 +++++++++++++-
 authentik/flows/views.py            |  2 +-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/authentik/flows/tests/test_views.py b/authentik/flows/tests/test_views.py
index 2afabcbe7..a89bea69e 100644
--- a/authentik/flows/tests/test_views.py
+++ b/authentik/flows/tests/test_views.py
@@ -8,7 +8,7 @@ from django.test.client import RequestFactory
 from django.utils.encoding import force_str
 
 from authentik.core.models import User
-from authentik.flows.exceptions import EmptyFlowException, FlowNonApplicableException
+from authentik.flows.exceptions import FlowNonApplicableException
 from authentik.flows.markers import ReevaluateMarker, StageMarker
 from authentik.flows.models import Flow, FlowDesignation, FlowStageBinding
 from authentik.flows.planner import FlowPlan, FlowPlanner
@@ -40,6 +40,10 @@ class TestFlowExecutor(TestCase):
     def setUp(self):
         self.request_factory = RequestFactory()
 
+    @patch(
+        "authentik.flows.views.to_stage_response",
+        TO_STAGE_RESPONSE_MOCK,
+    )
     def test_existing_plan_diff_flow(self):
         """Check that a plan for a different flow cancels the current plan"""
         flow = Flow.objects.create(
@@ -108,6 +112,10 @@ class TestFlowExecutor(TestCase):
         self.assertEqual(response.status_code, 302)
         self.assertEqual(response.url, reverse("authentik_core:shell"))
 
+    @patch(
+        "authentik.flows.views.to_stage_response",
+        TO_STAGE_RESPONSE_MOCK,
+    )
     def test_invalid_flow_redirect(self):
         """Tests that an invalid flow still redirects"""
         flow = Flow.objects.create(
@@ -157,6 +165,10 @@ class TestFlowExecutor(TestCase):
         plan: FlowPlan = session[SESSION_KEY_PLAN]
         self.assertEqual(len(plan.stages), 1)
 
+    @patch(
+        "authentik.flows.views.to_stage_response",
+        TO_STAGE_RESPONSE_MOCK,
+    )
     def test_reevaluate_remove_last(self):
         """Test planner with re-evaluate (last stage is removed)"""
         flow = Flow.objects.create(
diff --git a/authentik/flows/views.py b/authentik/flows/views.py
index 1e167adc6..597dad7bd 100644
--- a/authentik/flows/views.py
+++ b/authentik/flows/views.py
@@ -149,7 +149,7 @@ class FlowExecutorView(View):
             NEXT_ARG_NAME, "authentik_core:shell"
         )
         self.cancel()
-        return redirect_with_qs(next_param)
+        return to_stage_response(self.request, redirect_with_qs(next_param))
 
     def stage_ok(self) -> HttpResponse:
         """Callback called by stages upon successful completion.