From 3b2b3262d7287e74ccf63e0a83d0cd8d42bce599 Mon Sep 17 00:00:00 2001
From: Jens Langhammer <jens.langhammer@beryju.org>
Date: Sun, 27 Jun 2021 18:47:04 +0200
Subject: [PATCH] flows: add FlowStageBinding to flow plan instead of just
 stage

Signed-off-by: Jens Langhammer <jens.langhammer@beryju.org>
---
 authentik/flows/markers.py                    | 28 ++++++++-----
 authentik/flows/planner.py                    | 36 ++++++++--------
 authentik/flows/tests/test_planner.py         |  4 +-
 authentik/flows/tests/test_views.py           | 41 ++++++++++---------
 authentik/flows/views.py                      | 20 ++++++---
 authentik/providers/oauth2/views/authorize.py |  2 +-
 authentik/stages/captcha/tests.py             |  6 ++-
 authentik/stages/consent/tests.py             | 14 ++++---
 authentik/stages/deny/tests.py                |  6 ++-
 authentik/stages/email/tests/test_sending.py  |  8 ++--
 authentik/stages/email/tests/test_stage.py    | 12 +++---
 authentik/stages/invitation/tests.py          | 12 +++---
 authentik/stages/password/tests.py            | 16 ++++----
 authentik/stages/prompt/tests.py              | 16 ++++----
 authentik/stages/user_delete/tests.py         |  8 ++--
 authentik/stages/user_login/tests.py          | 10 +++--
 authentik/stages/user_logout/tests.py         |  6 ++-
 authentik/stages/user_write/tests.py          | 14 ++++---
 18 files changed, 151 insertions(+), 108 deletions(-)

diff --git a/authentik/flows/markers.py b/authentik/flows/markers.py
index e0b773982..d5bc297f1 100644
--- a/authentik/flows/markers.py
+++ b/authentik/flows/markers.py
@@ -6,7 +6,7 @@ from django.http.request import HttpRequest
 from structlog.stdlib import get_logger
 
 from authentik.core.models import User
-from authentik.flows.models import Stage
+from authentik.flows.models import FlowStageBinding
 from authentik.policies.engine import PolicyEngine
 from authentik.policies.models import PolicyBinding
 
@@ -22,11 +22,14 @@ class StageMarker:
 
     # pylint: disable=unused-argument
     def process(
-        self, plan: "FlowPlan", stage: Stage, http_request: Optional[HttpRequest]
-    ) -> Optional[Stage]:
+        self,
+        plan: "FlowPlan",
+        binding: FlowStageBinding,
+        http_request: Optional[HttpRequest],
+    ) -> Optional[FlowStageBinding]:
         """Process callback for this marker. This should be overridden by sub-classes.
         If a stage should be removed, return None."""
-        return stage
+        return binding
 
 
 @dataclass
@@ -37,13 +40,16 @@ class ReevaluateMarker(StageMarker):
     user: User
 
     def process(
-        self, plan: "FlowPlan", stage: Stage, http_request: Optional[HttpRequest]
-    ) -> Optional[Stage]:
+        self,
+        plan: "FlowPlan",
+        binding: FlowStageBinding,
+        http_request: Optional[HttpRequest],
+    ) -> Optional[FlowStageBinding]:
         """Re-evaluate policies bound to stage, and if they fail, remove from plan"""
         LOGGER.debug(
             "f(plan_inst)[re-eval marker]: running re-evaluation",
-            stage=stage,
-            binding=self.binding,
+            binding=binding,
+            policy_binding=self.binding,
         )
         engine = PolicyEngine(self.binding, self.user)
         engine.use_cache = False
@@ -53,10 +59,10 @@ class ReevaluateMarker(StageMarker):
         engine.build()
         result = engine.result
         if result.passing:
-            return stage
+            return binding
         LOGGER.warning(
-            "f(plan_inst)[re-eval marker]: stage failed re-evaluation",
-            stage=stage,
+            "f(plan_inst)[re-eval marker]: binding failed re-evaluation",
+            binding=binding,
             messages=result.messages,
         )
         return None
diff --git a/authentik/flows/planner.py b/authentik/flows/planner.py
index 9edb719fa..9f4174a12 100644
--- a/authentik/flows/planner.py
+++ b/authentik/flows/planner.py
@@ -13,7 +13,7 @@ from authentik.core.models import User
 from authentik.events.models import cleanse_dict
 from authentik.flows.exceptions import EmptyFlowException, FlowNonApplicableException
 from authentik.flows.markers import ReevaluateMarker, StageMarker
-from authentik.flows.models import Flow, FlowStageBinding, Stage
+from authentik.flows.models import Flow, FlowStageBinding
 from authentik.lib.config import CONFIG
 from authentik.policies.engine import PolicyEngine
 from authentik.root.monitoring import UpdatingGauge
@@ -52,33 +52,37 @@ class FlowPlan:
 
     flow_pk: str
 
-    stages: list[Stage] = field(default_factory=list)
+    bindings: list[FlowStageBinding] = field(default_factory=list)
     context: dict[str, Any] = field(default_factory=dict)
     markers: list[StageMarker] = field(default_factory=list)
 
-    def append(self, stage: Stage, marker: Optional[StageMarker] = None):
+    def append(self, binding: FlowStageBinding, marker: Optional[StageMarker] = None):
         """Append `stage` to all stages, optionall with stage marker"""
-        self.stages.append(stage)
+        self.bindings.append(binding)
         self.markers.append(marker or StageMarker())
 
-    def insert(self, stage: Stage, marker: Optional[StageMarker] = None):
+    def insert(self, binding: FlowStageBinding, marker: Optional[StageMarker] = None):
         """Insert stage into plan, as immediate next stage"""
-        self.stages.insert(1, stage)
+        self.bindings.insert(1, binding)
         self.markers.insert(1, marker or StageMarker())
 
-    def next(self, http_request: Optional[HttpRequest]) -> Optional[Stage]:
+    def next(self, http_request: Optional[HttpRequest]) -> Optional[FlowStageBinding]:
         """Return next pending stage from the bottom of the list"""
         if not self.has_stages:
             return None
-        stage = self.stages[0]
+        binding = self.bindings[0]
         marker = self.markers[0]
 
         if marker.__class__ is not StageMarker:
-            LOGGER.debug("f(plan_inst): stage has marker", stage=stage, marker=marker)
-        marked_stage = marker.process(self, stage, http_request)
+            LOGGER.debug(
+                "f(plan_inst): stage has marker", binding=binding, marker=marker
+            )
+        marked_stage = marker.process(self, binding, http_request)
         if not marked_stage:
-            LOGGER.debug("f(plan_inst): marker returned none, next stage", stage=stage)
-            self.stages.remove(stage)
+            LOGGER.debug(
+                "f(plan_inst): marker returned none, next stage", binding=binding
+            )
+            self.bindings.remove(binding)
             self.markers.remove(marker)
             if not self.has_stages:
                 return None
@@ -89,12 +93,12 @@ class FlowPlan:
     def pop(self):
         """Pop next pending stage from bottom of list"""
         self.markers.pop(0)
-        self.stages.pop(0)
+        self.bindings.pop(0)
 
     @property
     def has_stages(self) -> bool:
         """Check if there are any stages left in this plan"""
-        return len(self.markers) + len(self.stages) > 0
+        return len(self.markers) + len(self.bindings) > 0
 
 
 class FlowPlanner:
@@ -161,7 +165,7 @@ class FlowPlanner:
             plan = self._build_plan(user, request, default_context)
             cache.set(cache_key(self.flow, user), plan, CACHE_TIMEOUT)
             GAUGE_FLOWS_CACHED.update()
-            if not plan.stages and not self.allow_empty_flows:
+            if not plan.bindings and not self.allow_empty_flows:
                 raise EmptyFlowException()
             return plan
 
@@ -218,7 +222,7 @@ class FlowPlanner:
                     )
                     marker = ReevaluateMarker(binding=binding, user=user)
                 if stage:
-                    plan.append(stage, marker)
+                    plan.append(binding, marker)
             HIST_FLOWS_PLAN_TIME.labels(flow_slug=self.flow.slug)
         self._logger.debug(
             "f(plan): finished building",
diff --git a/authentik/flows/tests/test_planner.py b/authentik/flows/tests/test_planner.py
index 9eb60a112..8e185da58 100644
--- a/authentik/flows/tests/test_planner.py
+++ b/authentik/flows/tests/test_planner.py
@@ -182,8 +182,8 @@ class TestFlowPlanner(TestCase):
             planner = FlowPlanner(flow)
             plan = planner.plan(request)
 
-            self.assertEqual(plan.stages[0], binding.stage)
-            self.assertEqual(plan.stages[1], binding2.stage)
+            self.assertEqual(plan.bindings[0], binding)
+            self.assertEqual(plan.bindings[1], binding2)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], ReevaluateMarker)
diff --git a/authentik/flows/tests/test_views.py b/authentik/flows/tests/test_views.py
index fee5ff259..6b25ff970 100644
--- a/authentik/flows/tests/test_views.py
+++ b/authentik/flows/tests/test_views.py
@@ -52,8 +52,9 @@ class TestFlowExecutor(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         stage = DummyStage.objects.create(name="dummy")
+        binding = FlowStageBinding.objects.create(target=flow, stage=stage)
         plan = FlowPlan(
-            flow_pk=flow.pk.hex + "a", stages=[stage], markers=[StageMarker()]
+            flow_pk=flow.pk.hex + "a", bindings=[binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -163,7 +164,7 @@ class TestFlowExecutor(TestCase):
         # Check that two stages are in plan
         session = self.client.session
         plan: FlowPlan = session[SESSION_KEY_PLAN]
-        self.assertEqual(len(plan.stages), 2)
+        self.assertEqual(len(plan.bindings), 2)
         # Second request, submit form, one stage left
         response = self.client.post(exec_url)
         # Second request redirects to the same URL
@@ -172,7 +173,7 @@ class TestFlowExecutor(TestCase):
         # Check that two stages are in plan
         session = self.client.session
         plan: FlowPlan = session[SESSION_KEY_PLAN]
-        self.assertEqual(len(plan.stages), 1)
+        self.assertEqual(len(plan.bindings), 1)
 
     @patch(
         "authentik.flows.views.to_stage_response",
@@ -213,8 +214,8 @@ class TestFlowExecutor(TestCase):
 
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding.stage)
-            self.assertEqual(plan.stages[1], binding2.stage)
+            self.assertEqual(plan.bindings[0], binding)
+            self.assertEqual(plan.bindings[1], binding2)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], ReevaluateMarker)
@@ -267,9 +268,9 @@ class TestFlowExecutor(TestCase):
             self.assertEqual(response.status_code, 200)
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding.stage)
-            self.assertEqual(plan.stages[1], binding2.stage)
-            self.assertEqual(plan.stages[2], binding3.stage)
+            self.assertEqual(plan.bindings[0], binding)
+            self.assertEqual(plan.bindings[1], binding2)
+            self.assertEqual(plan.bindings[2], binding3)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], ReevaluateMarker)
@@ -281,8 +282,8 @@ class TestFlowExecutor(TestCase):
 
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding2.stage)
-            self.assertEqual(plan.stages[1], binding3.stage)
+            self.assertEqual(plan.bindings[0], binding2)
+            self.assertEqual(plan.bindings[1], binding3)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], StageMarker)
@@ -338,9 +339,9 @@ class TestFlowExecutor(TestCase):
             self.assertEqual(response.status_code, 200)
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding.stage)
-            self.assertEqual(plan.stages[1], binding2.stage)
-            self.assertEqual(plan.stages[2], binding3.stage)
+            self.assertEqual(plan.bindings[0], binding)
+            self.assertEqual(plan.bindings[1], binding2)
+            self.assertEqual(plan.bindings[2], binding3)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], ReevaluateMarker)
@@ -352,8 +353,8 @@ class TestFlowExecutor(TestCase):
 
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding2.stage)
-            self.assertEqual(plan.stages[1], binding3.stage)
+            self.assertEqual(plan.bindings[0], binding2)
+            self.assertEqual(plan.bindings[1], binding3)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], StageMarker)
@@ -364,7 +365,7 @@ class TestFlowExecutor(TestCase):
 
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding3.stage)
+            self.assertEqual(plan.bindings[0], binding3)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
 
@@ -438,10 +439,10 @@ class TestFlowExecutor(TestCase):
 
             plan: FlowPlan = self.client.session[SESSION_KEY_PLAN]
 
-            self.assertEqual(plan.stages[0], binding.stage)
-            self.assertEqual(plan.stages[1], binding2.stage)
-            self.assertEqual(plan.stages[2], binding3.stage)
-            self.assertEqual(plan.stages[3], binding4.stage)
+            self.assertEqual(plan.bindings[0], binding)
+            self.assertEqual(plan.bindings[1], binding2)
+            self.assertEqual(plan.bindings[2], binding3)
+            self.assertEqual(plan.bindings[3], binding4)
 
             self.assertIsInstance(plan.markers[0], StageMarker)
             self.assertIsInstance(plan.markers[1], ReevaluateMarker)
diff --git a/authentik/flows/views.py b/authentik/flows/views.py
index 499338d97..c527d1b16 100644
--- a/authentik/flows/views.py
+++ b/authentik/flows/views.py
@@ -37,7 +37,13 @@ from authentik.flows.challenge import (
     WithUserInfoChallenge,
 )
 from authentik.flows.exceptions import EmptyFlowException, FlowNonApplicableException
-from authentik.flows.models import ConfigurableStage, Flow, FlowDesignation, Stage
+from authentik.flows.models import (
+    ConfigurableStage,
+    Flow,
+    FlowDesignation,
+    FlowStageBinding,
+    Stage,
+)
 from authentik.flows.planner import (
     PLAN_CONTEXT_PENDING_USER,
     PLAN_CONTEXT_REDIRECT,
@@ -107,6 +113,7 @@ class FlowExecutorView(APIView):
     flow: Flow
 
     plan: Optional[FlowPlan] = None
+    current_binding: FlowStageBinding
     current_stage: Stage
     current_stage_view: View
 
@@ -159,11 +166,12 @@ class FlowExecutorView(APIView):
         request.session[SESSION_KEY_GET] = QueryDict(request.GET.get("query", ""))
         # We don't save the Plan after getting the next stage
         # as it hasn't been successfully passed yet
-        next_stage = self.plan.next(self.request)
-        if not next_stage:
+        next_binding = self.plan.next(self.request)
+        if not next_binding:
             self._logger.debug("f(exec): no more stages, flow is done.")
             return self._flow_done()
-        self.current_stage = next_stage
+        self.current_binding = next_binding
+        self.current_stage = next_binding.stage
         self._logger.debug(
             "f(exec): Current stage",
             current_stage=self.current_stage,
@@ -293,10 +301,10 @@ class FlowExecutorView(APIView):
         )
         self.plan.pop()
         self.request.session[SESSION_KEY_PLAN] = self.plan
-        if self.plan.stages:
+        if self.plan.bindings:
             self._logger.debug(
                 "f(exec): Continuing with next stage",
-                remaining=len(self.plan.stages),
+                remaining=len(self.plan.bindings),
             )
             kwargs = self.kwargs
             kwargs.update({"flow_slug": self.flow.slug})
diff --git a/authentik/providers/oauth2/views/authorize.py b/authentik/providers/oauth2/views/authorize.py
index 3728c92b5..1d0704c5d 100644
--- a/authentik/providers/oauth2/views/authorize.py
+++ b/authentik/providers/oauth2/views/authorize.py
@@ -468,7 +468,7 @@ class AuthorizationFlowInitView(PolicyAccessView):
         # OpenID clients can specify a `prompt` parameter, and if its set to consent we
         # need to inject a consent stage
         if PROMPT_CONSNET in self.params.prompt:
-            if not any(isinstance(x, ConsentStageView) for x in plan.stages):
+            if not any(isinstance(x.stage, ConsentStageView) for x in plan.bindings):
                 # Plan does not have any consent stage, so we add an in-memory one
                 stage = ConsentStage(
                     name="OAuth2 Provider In-memory consent stage",
diff --git a/authentik/stages/captcha/tests.py b/authentik/stages/captcha/tests.py
index 3579438e5..8c863e2ee 100644
--- a/authentik/stages/captcha/tests.py
+++ b/authentik/stages/captcha/tests.py
@@ -36,12 +36,14 @@ class TestCaptchaStage(TestCase):
             public_key=RECAPTCHA_PUBLIC_KEY,
             private_key=RECAPTCHA_PRIVATE_KEY,
         )
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_valid(self):
         """Test valid captcha"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
diff --git a/authentik/stages/consent/tests.py b/authentik/stages/consent/tests.py
index d395af6a0..728b61649 100644
--- a/authentik/stages/consent/tests.py
+++ b/authentik/stages/consent/tests.py
@@ -39,9 +39,11 @@ class TestConsentStage(TestCase):
         stage = ConsentStage.objects.create(
             name="consent", mode=ConsentMode.ALWAYS_REQUIRE
         )
-        FlowStageBinding.objects.create(target=flow, stage=stage, order=2)
+        binding = FlowStageBinding.objects.create(target=flow, stage=stage, order=2)
 
-        plan = FlowPlan(flow_pk=flow.pk.hex, stages=[stage], markers=[StageMarker()])
+        plan = FlowPlan(
+            flow_pk=flow.pk.hex, bindings=[binding], markers=[StageMarker()]
+        )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
         session.save()
@@ -69,11 +71,11 @@ class TestConsentStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         stage = ConsentStage.objects.create(name="consent", mode=ConsentMode.PERMANENT)
-        FlowStageBinding.objects.create(target=flow, stage=stage, order=2)
+        binding = FlowStageBinding.objects.create(target=flow, stage=stage, order=2)
 
         plan = FlowPlan(
             flow_pk=flow.pk.hex,
-            stages=[stage],
+            bindings=[binding],
             markers=[StageMarker()],
             context={PLAN_CONTEXT_APPLICATION: self.application},
         )
@@ -110,11 +112,11 @@ class TestConsentStage(TestCase):
         stage = ConsentStage.objects.create(
             name="consent", mode=ConsentMode.EXPIRING, consent_expire_in="seconds=1"
         )
-        FlowStageBinding.objects.create(target=flow, stage=stage, order=2)
+        binding = FlowStageBinding.objects.create(target=flow, stage=stage, order=2)
 
         plan = FlowPlan(
             flow_pk=flow.pk.hex,
-            stages=[stage],
+            bindings=[binding],
             markers=[StageMarker()],
             context={PLAN_CONTEXT_APPLICATION: self.application},
         )
diff --git a/authentik/stages/deny/tests.py b/authentik/stages/deny/tests.py
index 0df3d9bfd..9a15181ce 100644
--- a/authentik/stages/deny/tests.py
+++ b/authentik/stages/deny/tests.py
@@ -26,12 +26,14 @@ class TestUserDenyStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.stage = DenyStage.objects.create(name="logout")
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_valid_password(self):
         """Test with a valid pending user and backend"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
diff --git a/authentik/stages/email/tests/test_sending.py b/authentik/stages/email/tests/test_sending.py
index 9b12201f5..5467999b4 100644
--- a/authentik/stages/email/tests/test_sending.py
+++ b/authentik/stages/email/tests/test_sending.py
@@ -34,12 +34,14 @@ class TestEmailStageSending(TestCase):
         self.stage = EmailStage.objects.create(
             name="email",
         )
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_pending_user(self):
         """Test with pending user"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -67,7 +69,7 @@ class TestEmailStageSending(TestCase):
     def test_send_error(self):
         """Test error during sending (sending will be retried)"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
diff --git a/authentik/stages/email/tests/test_stage.py b/authentik/stages/email/tests/test_stage.py
index ae499b05b..541e21750 100644
--- a/authentik/stages/email/tests/test_stage.py
+++ b/authentik/stages/email/tests/test_stage.py
@@ -35,12 +35,14 @@ class TestEmailStage(TestCase):
         self.stage = EmailStage.objects.create(
             name="email",
         )
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_rendering(self):
         """Test with pending user"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -56,7 +58,7 @@ class TestEmailStage(TestCase):
     def test_without_user(self):
         """Test without pending user"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -71,7 +73,7 @@ class TestEmailStage(TestCase):
     def test_pending_user(self):
         """Test with pending user"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -102,7 +104,7 @@ class TestEmailStage(TestCase):
         # Make sure token exists
         self.test_pending_user()
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
diff --git a/authentik/stages/invitation/tests.py b/authentik/stages/invitation/tests.py
index 509e5fcb8..72e9f93bf 100644
--- a/authentik/stages/invitation/tests.py
+++ b/authentik/stages/invitation/tests.py
@@ -35,7 +35,9 @@ class TestUserLoginStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.stage = InvitationStage.objects.create(name="invitation")
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     @patch(
         "authentik.flows.views.to_stage_response",
@@ -44,7 +46,7 @@ class TestUserLoginStage(TestCase):
     def test_without_invitation_fail(self):
         """Test without any invitation, continue_flow_without_invitation not set."""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         plan.context[PLAN_CONTEXT_AUTHENTICATION_BACKEND] = BACKEND_DJANGO
@@ -75,7 +77,7 @@ class TestUserLoginStage(TestCase):
         self.stage.continue_flow_without_invitation = True
         self.stage.save()
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         plan.context[PLAN_CONTEXT_AUTHENTICATION_BACKEND] = BACKEND_DJANGO
@@ -103,7 +105,7 @@ class TestUserLoginStage(TestCase):
     def test_with_invitation_get(self):
         """Test with invitation, check data in session"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -143,7 +145,7 @@ class TestUserLoginStage(TestCase):
         )
 
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PROMPT] = {INVITATION_TOKEN_KEY: invite.pk.hex}
         session = self.client.session
diff --git a/authentik/stages/password/tests.py b/authentik/stages/password/tests.py
index 846360421..5a653451f 100644
--- a/authentik/stages/password/tests.py
+++ b/authentik/stages/password/tests.py
@@ -39,7 +39,9 @@ class TestPasswordStage(TestCase):
         self.stage = PasswordStage.objects.create(
             name="password", backends=[BACKEND_DJANGO]
         )
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     @patch(
         "authentik.flows.views.to_stage_response",
@@ -48,7 +50,7 @@ class TestPasswordStage(TestCase):
     def test_without_user(self):
         """Test without user"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -84,7 +86,7 @@ class TestPasswordStage(TestCase):
         )
 
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -101,7 +103,7 @@ class TestPasswordStage(TestCase):
     def test_valid_password(self):
         """Test with a valid pending user and valid password"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -129,7 +131,7 @@ class TestPasswordStage(TestCase):
     def test_invalid_password(self):
         """Test with a valid pending user and invalid password"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -148,7 +150,7 @@ class TestPasswordStage(TestCase):
     def test_invalid_password_lockout(self):
         """Test with a valid pending user and invalid password (trigger logout counter)"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -189,7 +191,7 @@ class TestPasswordStage(TestCase):
         """Test with a valid pending user and valid password.
         Backend is patched to return PermissionError"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
diff --git a/authentik/stages/prompt/tests.py b/authentik/stages/prompt/tests.py
index 125b54bcc..cf485a2ef 100644
--- a/authentik/stages/prompt/tests.py
+++ b/authentik/stages/prompt/tests.py
@@ -102,12 +102,14 @@ class TestPromptStage(TestCase):
             hidden_prompt.field_key: hidden_prompt.placeholder,
         }
 
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_render(self):
         """Test render of form, check if all prompts are rendered correctly"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -125,7 +127,7 @@ class TestPromptStage(TestCase):
     def test_valid_challenge_with_policy(self) -> PromptChallengeResponse:
         """Test challenge_response validation"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         expr = "return request.context['password_prompt'] == request.context['password2_prompt']"
         expr_policy = ExpressionPolicy.objects.create(
@@ -142,7 +144,7 @@ class TestPromptStage(TestCase):
     def test_invalid_challenge(self) -> PromptChallengeResponse:
         """Test challenge_response validation"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         expr = "False"
         expr_policy = ExpressionPolicy.objects.create(
@@ -159,7 +161,7 @@ class TestPromptStage(TestCase):
     def test_valid_challenge_request(self):
         """Test a request with valid challenge_response data"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -196,7 +198,7 @@ class TestPromptStage(TestCase):
     def test_invalid_password(self):
         """Test challenge_response validation"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         self.prompt_data["password2_prompt"] = "qwerqwerqr"
         challenge_response = PromptChallengeResponse(
@@ -215,7 +217,7 @@ class TestPromptStage(TestCase):
     def test_invalid_username(self):
         """Test challenge_response validation"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         self.prompt_data["username_prompt"] = "akadmin"
         challenge_response = PromptChallengeResponse(
diff --git a/authentik/stages/user_delete/tests.py b/authentik/stages/user_delete/tests.py
index 897194398..e1e357d61 100644
--- a/authentik/stages/user_delete/tests.py
+++ b/authentik/stages/user_delete/tests.py
@@ -30,7 +30,9 @@ class TestUserDeleteStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.stage = UserDeleteStage.objects.create(name="delete")
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     @patch(
         "authentik.flows.views.to_stage_response",
@@ -39,7 +41,7 @@ class TestUserDeleteStage(TestCase):
     def test_no_user(self):
         """Test without user set"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -66,7 +68,7 @@ class TestUserDeleteStage(TestCase):
     def test_user_delete_get(self):
         """Test Form render"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
diff --git a/authentik/stages/user_login/tests.py b/authentik/stages/user_login/tests.py
index f5538e9b7..ebcb90569 100644
--- a/authentik/stages/user_login/tests.py
+++ b/authentik/stages/user_login/tests.py
@@ -30,12 +30,14 @@ class TestUserLoginStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.stage = UserLoginStage.objects.create(name="login")
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_valid_password(self):
         """Test with a valid pending user and backend"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -61,7 +63,7 @@ class TestUserLoginStage(TestCase):
         self.stage.session_duration = "seconds=2"
         self.stage.save()
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         session = self.client.session
@@ -92,7 +94,7 @@ class TestUserLoginStage(TestCase):
     def test_without_user(self):
         """Test a plan without any pending user, resulting in a denied"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
diff --git a/authentik/stages/user_logout/tests.py b/authentik/stages/user_logout/tests.py
index 9f706ba5d..2472fc7cc 100644
--- a/authentik/stages/user_logout/tests.py
+++ b/authentik/stages/user_logout/tests.py
@@ -28,12 +28,14 @@ class TestUserLogoutStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.stage = UserLogoutStage.objects.create(name="logout")
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
 
     def test_valid_password(self):
         """Test with a valid pending user and backend"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = self.user
         plan.context[PLAN_CONTEXT_AUTHENTICATION_BACKEND] = BACKEND_DJANGO
diff --git a/authentik/stages/user_write/tests.py b/authentik/stages/user_write/tests.py
index 4d3b2767d..bd8fea12c 100644
--- a/authentik/stages/user_write/tests.py
+++ b/authentik/stages/user_write/tests.py
@@ -37,7 +37,9 @@ class TestUserWriteStage(TestCase):
             designation=FlowDesignation.AUTHENTICATION,
         )
         self.stage = UserWriteStage.objects.create(name="write")
-        FlowStageBinding.objects.create(target=self.flow, stage=self.stage, order=2)
+        self.binding = FlowStageBinding.objects.create(
+            target=self.flow, stage=self.stage, order=2
+        )
         self.source = Source.objects.create(name="fake_source")
 
     def test_user_create(self):
@@ -48,7 +50,7 @@ class TestUserWriteStage(TestCase):
         )
 
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PROMPT] = {
             "username": "test-user",
@@ -92,7 +94,7 @@ class TestUserWriteStage(TestCase):
             for _ in range(8)
         )
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         plan.context[PLAN_CONTEXT_PENDING_USER] = User.objects.create(
             username="unittest", email="test@beryju.org"
@@ -135,7 +137,7 @@ class TestUserWriteStage(TestCase):
     def test_without_data(self):
         """Test without data results in error"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         session[SESSION_KEY_PLAN] = plan
@@ -167,7 +169,7 @@ class TestUserWriteStage(TestCase):
     def test_blank_username(self):
         """Test with blank username results in error"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         plan.context[PLAN_CONTEXT_PROMPT] = {
@@ -204,7 +206,7 @@ class TestUserWriteStage(TestCase):
     def test_duplicate_data(self):
         """Test with duplicate data, should trigger error"""
         plan = FlowPlan(
-            flow_pk=self.flow.pk.hex, stages=[self.stage], markers=[StageMarker()]
+            flow_pk=self.flow.pk.hex, bindings=[self.binding], markers=[StageMarker()]
         )
         session = self.client.session
         plan.context[PLAN_CONTEXT_PROMPT] = {