tests/e2e: replace apply_default_data with @apply_migration decorator

tests/e2e/test_flows_authenticators.py

@@ -17,7 +17,7 @@ from authentik.flows.models import Flow, FlowStageBinding
 from authentik.stages.authenticator_static.models import AuthenticatorStaticStage
 from authentik.stages.authenticator_totp.models import AuthenticatorTOTPStage
 from authentik.stages.authenticator_validate.models import AuthenticatorValidateStage
-from tests.e2e.utils import USER, SeleniumTestCase, retry
+from tests.e2e.utils import USER, SeleniumTestCase, apply_migration, retry
 
 
 @skipUnless(platform.startswith("linux"), "requires local docker")
@@ -25,6 +25,8 @@ class TestFlowsAuthenticator(SeleniumTestCase):
     """test flow with otp stages"""
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
     def test_totp_validate(self):
         """test flow with otp stages"""
         sleep(1)
@@ -61,6 +63,9 @@ class TestFlowsAuthenticator(SeleniumTestCase):
         self.assert_user(USER())
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
+    @apply_migration("authentik_stages_authenticator_totp", "0006_default_setup_flow")
     def test_totp_setup(self):
         """test TOTP Setup stage"""
         flow: Flow = Flow.objects.get(slug="default-authentication-flow")
@@ -108,6 +113,9 @@ class TestFlowsAuthenticator(SeleniumTestCase):
         self.assertTrue(TOTPDevice.objects.filter(user=USER(), confirmed=True).exists())
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
+    @apply_migration("authentik_stages_authenticator_static", "0005_default_setup_flow")
     def test_static_setup(self):
         """test Static OTP Setup stage"""
         flow: Flow = Flow.objects.get(slug="default-authentication-flow")

tests/e2e/test_flows_enroll.py

@@ -16,7 +16,7 @@ from authentik.stages.identification.models import IdentificationStage
 from authentik.stages.prompt.models import FieldTypes, Prompt, PromptStage
 from authentik.stages.user_login.models import UserLoginStage
 from authentik.stages.user_write.models import UserWriteStage
-from tests.e2e.utils import USER, SeleniumTestCase, retry
+from tests.e2e.utils import USER, SeleniumTestCase, apply_migration, retry
 
 
 @skipUnless(platform.startswith("linux"), "requires local docker")
@@ -37,6 +37,8 @@ class TestFlowsEnroll(SeleniumTestCase):
         }
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
     # pylint: disable=too-many-locals
     def test_enroll_2_step(self):
         """Test 2-step enroll flow"""
@@ -101,6 +103,8 @@ class TestFlowsEnroll(SeleniumTestCase):
         self.assertEqual(user.email, "foo@bar.baz")
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
     @override_settings(EMAIL_BACKEND="django.core.mail.backends.smtp.EmailBackend")
     def test_enroll_email(self):
         """Test enroll with Email verification"""

tests/e2e/test_flows_login.py

@@ -2,7 +2,7 @@
 from sys import platform
 from unittest.case import skipUnless
 
-from tests.e2e.utils import USER, SeleniumTestCase, retry
+from tests.e2e.utils import USER, SeleniumTestCase, apply_migration, retry
 
 
 @skipUnless(platform.startswith("linux"), "requires local docker")
@@ -10,6 +10,8 @@ class TestFlowsLogin(SeleniumTestCase):
     """test default login flow"""
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
     def test_login(self):
         """test default login flow"""
         self.driver.get(f"{self.live_server_url}/flows/default-authentication-flow/")

tests/e2e/test_flows_stage_setup.py

@@ -1,5 +1,6 @@
 """test stage setup flows (password change)"""
 from sys import platform
+from time import sleep
 from unittest.case import skipUnless
 
 from selenium.webdriver.common.by import By
@@ -9,7 +10,7 @@ from authentik.core.models import User
 from authentik.flows.models import Flow, FlowDesignation
 from authentik.providers.oauth2.generators import generate_client_secret
 from authentik.stages.password.models import PasswordStage
-from tests.e2e.utils import USER, SeleniumTestCase, retry
+from tests.e2e.utils import USER, SeleniumTestCase, apply_migration, retry
 
 
 @skipUnless(platform.startswith("linux"), "requires local docker")
@@ -17,6 +18,9 @@ class TestFlowsStageSetup(SeleniumTestCase):
     """test stage setup flows"""
 
     @retry()
+    @apply_migration("authentik_core", "0003_default_user")
+    @apply_migration("authentik_flows", "0008_default_flows")
+    @apply_migration("authentik_stages_password", "0002_passwordstage_change_flow")
     def test_password_change(self):
         """test password change flow"""
         # Ensure that password stage has change_flow set
@@ -34,10 +38,7 @@ class TestFlowsStageSetup(SeleniumTestCase):
         self.driver.get(
             f"{self.live_server_url}/flows/default-authentication-flow/?next=%2F"
         )
-        self.driver.find_element(By.ID, "id_uid_field").send_keys(USER().username)
-        self.driver.find_element(By.ID, "id_uid_field").send_keys(Keys.ENTER)
-        self.driver.find_element(By.ID, "id_password").send_keys(USER().username)
-        self.driver.find_element(By.ID, "id_password").send_keys(Keys.ENTER)
+        self.login()
         self.wait_for_url(self.shell_url("/library"))
 
         self.driver.get(
@@ -46,10 +47,19 @@ class TestFlowsStageSetup(SeleniumTestCase):
                 stage_uuid=PasswordStage.objects.first().stage_uuid,
             )
         )
-        self.driver.find_element(By.ID, "id_password").send_keys(new_password)
-        self.driver.find_element(By.ID, "id_password_repeat").click()
-        self.driver.find_element(By.ID, "id_password_repeat").send_keys(new_password)
-        self.driver.find_element(By.CSS_SELECTOR, ".pf-c-button").click()
+
+        flow_executor = self.get_shadow_root("ak-flow-executor")
+        prompt_stage = self.get_shadow_root("ak-stage-prompt", flow_executor)
+
+        prompt_stage.find_element(By.CSS_SELECTOR, "input[name=password]").send_keys(
+            new_password
+        )
+        prompt_stage.find_element(
+            By.CSS_SELECTOR, "input[name=password_repeat]"
+        ).send_keys(new_password)
+        prompt_stage.find_element(
+            By.CSS_SELECTOR, "input[name=password_repeat]"
+        ).send_keys(Keys.ENTER)
 
         self.wait_for_url(self.shell_url("/library"))
         # Because USER() is cached, we need to get the user manually here

tests/e2e/utils.py

@@ -1,6 +1,6 @@
 """authentik e2e testing utilities"""
 import json
-from functools import wraps
+from functools import lru_cache, wraps
 from glob import glob
 from importlib.util import module_from_spec, spec_from_file_location
 from inspect import getmembers, isfunction
@@ -57,7 +57,6 @@ class SeleniumTestCase(StaticLiveServerTestCase):
         self.driver.maximize_window()
         self.driver.implicitly_wait(30)
         self.wait = WebDriverWait(self.driver, self.wait_timeout)
-        self.apply_default_data()
         self.logger = get_logger()
         if specs := self.get_container_specs():
             self.container = self._start_container(specs)
@@ -166,35 +165,12 @@ class SeleniumTestCase(StaticLiveServerTestCase):
         self.assertEqual(user["name"].value, expected_user.name)
         self.assertEqual(user["email"].value, expected_user.email)
 
-    def apply_default_data(self):
-        """apply objects created by migrations after tables have been truncated"""
-        # Not all default objects are managed, like users for example
-        # Hence we still have to load all migrations and apply them, then run the ObjectManager
-        # Find all migration files
-        # load all functions
-        migration_files = glob("**/migrations/*.py", recursive=True)
-        matches = []
-        for migration in migration_files:
-            with open(migration, "r+") as migration_file:
-                # Check if they have a `RunPython`
-                if "RunPython" in migration_file.read():
-                    matches.append(migration)
 
-        with connection.schema_editor() as schema_editor:
-            for match in matches:
-                # Load module from file path
-                spec = spec_from_file_location("", match)
-                migration_module = module_from_spec(spec)
-                # pyright: reportGeneralTypeIssues=false
-                spec.loader.exec_module(migration_module)
-                # Call all functions from module
-                for _, func in getmembers(migration_module, isfunction):
-                    with transaction.atomic():
-                        try:
-                            func(apps, schema_editor)
-                        except IntegrityError:
-                            pass
-        ObjectManager().run()
+@lru_cache
+def get_loader():
+    """Thin wrapper to lazily get a Migration Loader, only when it's needed
+    and only once"""
+    return MigrationLoader(connection)
 
 
 def apply_migration(app_name: str, migration_name: str):
@@ -203,11 +179,9 @@ def apply_migration(app_name: str, migration_name: str):
     def wrapper_outter(func: Callable):
         """Retry test multiple times"""
 
-        loader = MigrationLoader(connection)
-
         @wraps(func)
         def wrapper(self: TransactionTestCase, *args, **kwargs):
-            migration = loader.get_migration(app_name, migration_name)
+            migration = get_loader().get_migration(app_name, migration_name)
             with connection.schema_editor() as schema_editor:
                 for operation in migration.operations:
                     if not isinstance(operation, RunPython):