enterprise/stages/source: fix dispatch method signature (cherry-pick #13321) (#13326)

enterprise/stages/source: fix dispatch method signature (#13321)

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
Co-authored-by: Jens L. <jens@goauthentik.io>

authentik/enterprise/stages/source/stage.py

@@ -89,9 +89,9 @@ class SourceStageFinal(StageView):
     This stage uses the override flow token to resume execution of the initial flow the
     source stage is bound to."""
 
-    def dispatch(self):
+    def dispatch(self, *args, **kwargs):
         token: FlowToken = self.request.session.get(SESSION_KEY_OVERRIDE_FLOW_TOKEN)
-        self._logger.info("Replacing source flow with overridden flow", flow=token.flow.slug)
+        self.logger.info("Replacing source flow with overridden flow", flow=token.flow.slug)
         plan = token.plan
         plan.context[PLAN_CONTEXT_IS_RESTORED] = token
         response = plan.to_redirect(self.request, token.flow)

authentik/enterprise/stages/source/tests.py

@@ -4,7 +4,8 @@ from django.urls import reverse
 
 from authentik.core.tests.utils import create_test_flow, create_test_user
 from authentik.enterprise.stages.source.models import SourceStage
-from authentik.flows.models import FlowDesignation, FlowStageBinding, FlowToken
+from authentik.enterprise.stages.source.stage import SourceStageFinal
+from authentik.flows.models import FlowDesignation, FlowStageBinding, FlowToken, in_memory_stage
 from authentik.flows.planner import PLAN_CONTEXT_IS_RESTORED, FlowPlan
 from authentik.flows.tests import FlowTestCase
 from authentik.flows.views.executor import SESSION_KEY_PLAN
@@ -87,6 +88,7 @@ class TestSourceStage(FlowTestCase):
         self.assertIsNotNone(flow_token)
         session = self.client.session
         plan: FlowPlan = session[SESSION_KEY_PLAN]
+        plan.insert_stage(in_memory_stage(SourceStageFinal), index=0)
         plan.context[PLAN_CONTEXT_IS_RESTORED] = flow_token
         session[SESSION_KEY_PLAN] = plan
         session.save()
@@ -96,4 +98,6 @@ class TestSourceStage(FlowTestCase):
             reverse("authentik_api:flow-executor", kwargs={"flow_slug": flow.slug}), follow=True
         )
         self.assertEqual(response.status_code, 200)
-        self.assertStageRedirects(response, reverse("authentik_core:root-redirect"))
+        self.assertStageRedirects(
+            response, reverse("authentik_core:if-flow", kwargs={"flow_slug": flow.slug})
+        )

authentik/flows/planner.py

@@ -76,10 +76,10 @@ class FlowPlan:
         self.bindings.append(binding)
         self.markers.append(marker or StageMarker())
 
-    def insert_stage(self, stage: Stage, marker: StageMarker | None = None):
+    def insert_stage(self, stage: Stage, marker: StageMarker | None = None, index=1):
         """Insert stage into plan, as immediate next stage"""
-        self.bindings.insert(1, FlowStageBinding(stage=stage, order=0))
-        self.markers.insert(1, marker or StageMarker())
+        self.bindings.insert(index, FlowStageBinding(stage=stage, order=0))
+        self.markers.insert(index, marker or StageMarker())
 
     def redirect(self, destination: str):
         """Insert a redirect stage as next stage"""