tests: refactor test harness to split apart a single file

Makefile

@@ -144,8 +144,14 @@ dev-create-db:
 dev-reset: dev-drop-db dev-create-db migrate  ## Drop and restore the Authentik PostgreSQL instance to a "fresh install" state.
 
 update-test-mmdb:  ## Update test GeoIP and ASN Databases
-	curl -L https://raw.githubusercontent.com/maxmind/MaxMind-DB/refs/heads/main/test-data/GeoLite2-ASN-Test.mmdb -o ${PWD}/tests/GeoLite2-ASN-Test.mmdb
-	curl -L https://raw.githubusercontent.com/maxmind/MaxMind-DB/refs/heads/main/test-data/GeoLite2-City-Test.mmdb -o ${PWD}/tests/GeoLite2-City-Test.mmdb
+	curl \
+		-L \
+		-o ${PWD}/tests/geoip/GeoLite2-ASN-Test.mmdb \
+		https://raw.githubusercontent.com/maxmind/MaxMind-DB/refs/heads/main/test-data/GeoLite2-ASN-Test.mmdb
+	curl \
+		-L \
+		-o ${PWD}/tests/geoip/GeoLite2-City-Test.mmdb \
+		https://raw.githubusercontent.com/maxmind/MaxMind-DB/refs/heads/main/test-data/GeoLite2-City-Test.mmdb
 
 bump:  ## Bump authentik version. Usage: make bump version=20xx.xx.xx
 ifndef version

authentik/root/test_plugin.py

@@ -6,7 +6,7 @@
 from cryptography.hazmat.backends.openssl.backend import backend
 
 from authentik import authentik_full_version
-from tests.e2e.utils import get_local_ip
+from tests.decorators import get_local_ip
 
 IS_CI = "CI" in environ
 

authentik/root/test_runner.py

@@ -69,8 +69,8 @@ def _setup_test_environment(self):
 
         # Test-specific configuration
         test_config = {
-            "events.context_processors.geoip": "tests/GeoLite2-City-Test.mmdb",
-            "events.context_processors.asn": "tests/GeoLite2-ASN-Test.mmdb",
+            "events.context_processors.geoip": "tests/geoip/GeoLite2-City-Test.mmdb",
+            "events.context_processors.asn": "tests/geoip/GeoLite2-ASN-Test.mmdb",
             "blueprints_dir": "./blueprints",
             "outposts.container_image_base": f"ghcr.io/goauthentik/dev-%(type)s:{get_docker_tag()}",
             "tenants.enabled": False,

internal/outpost/flow/session.go

@@ -15,5 +15,14 @@ func (fe *FlowExecutor) Session() *jwt.Token {
 		return nil
 	}
 	t, _, _ := jwt.NewParser().ParseUnverified(sc.Value, &SessionCookieClaims{})
+	// During testing the session cookie value is not a JWT but rather just the session ID
+	// in which case we wrap that in a pseudo-JWT
+	if t == nil {
+		return &jwt.Token{
+			Claims: &SessionCookieClaims{
+				SessionID: sc.Value,
+			},
+		}
+	}
 	return t
 }

internal/outpost/ldap/ldap.go

@@ -149,6 +149,7 @@ func (ls *LDAPServer) handleWSSessionEnd(ctx context.Context, msg ak.Event) erro
 	ls.log.Info("Disconnecting session due to session end event")
 	conn, ok := ls.connections[mmsg.SessionID]
 	if !ok {
+		ls.log.Warn("Could not disconnect session, connection not found")
 		return nil
 	}
 	delete(ls.connections, mmsg.SessionID)

pyproject.toml

@@ -256,6 +256,7 @@ module = [
   "authentik.tenants.*",
   "guardian.*",
   "lifecycle.*",
+  "tests.*",
   "tests.e2e.*",
   "tests.integration.*",
   "tests.openid_conformance.*",

scripts/generate_config.py

@@ -27,8 +27,8 @@ def generate_local_config() -> dict[str, Any]:
         "cert_discovery_dir": "./certs",
         "events": {
             "processors": {
-                "geoip": "tests/GeoLite2-City-Test.mmdb",
-                "asn": "tests/GeoLite2-ASN-Test.mmdb",
+                "geoip": "tests/geoip/GeoLite2-City-Test.mmdb",
+                "asn": "tests/geoip/GeoLite2-ASN-Test.mmdb",
             }
         },
         "storage": {

tests/decorators.py

@@ -0,0 +1,77 @@
+import socket
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from os import environ, getenv
+from typing import Any
+
+from django.db import connection
+from django.db.migrations.loader import MigrationLoader
+from django.test.testcases import TransactionTestCase
+from selenium.common.exceptions import (
+    NoSuchElementException,
+    TimeoutException,
+    WebDriverException,
+)
+from structlog.stdlib import get_logger
+
+IS_CI = "CI" in environ
+RETRIES = int(environ.get("RETRIES", "3")) if IS_CI else 1
+SHADOW_ROOT_RETRIES = 5
+
+JSONType = dict[str, Any] | list[Any] | str | int | float | bool | None
+
+
+def get_local_ip(override=True) -> str:
+    """Get the local machine's IP"""
+    if (local_ip := getenv("LOCAL_IP")) and override:
+        return local_ip
+    hostname = socket.gethostname()
+    try:
+        return socket.gethostbyname(hostname)
+    except socket.gaierror:
+        return "0.0.0.0"
+
+
+@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 retry(max_retires=RETRIES, exceptions=None):
+    """Retry test multiple times. Default to catching Selenium Timeout Exception"""
+
+    if not exceptions:
+        exceptions = [WebDriverException, TimeoutException, NoSuchElementException]
+
+    logger = get_logger()
+
+    def retry_actual(func: Callable):
+        """Retry test multiple times"""
+        count = 1
+
+        @wraps(func)
+        def wrapper(self: TransactionTestCase, *args, **kwargs):
+            """Run test again if we're below max_retries, including tearDown and
+            setUp. Otherwise raise the error"""
+            nonlocal count
+            try:
+                return func(self, *args, **kwargs)
+
+            except tuple(exceptions) as exc:
+                count += 1
+                if count > max_retires:
+                    logger.debug("Exceeded retry count", exc=exc, test=self)
+
+                    raise exc
+                logger.debug("Retrying on error", exc=exc, test=self)
+                self.tearDown()
+                self._post_teardown()
+                self._pre_setup()
+                self.setUp()
+                return wrapper(self, *args, **kwargs)
+
+        return wrapper
+
+    return retry_actual

tests/docker.py

@@ -1,6 +1,5 @@
 """authentik e2e testing utilities"""
 
-from os import environ
 from time import sleep
 from typing import Any
 from unittest.case import TestCase
@@ -13,8 +12,6 @@
 from authentik.lib.generators import generate_id
 from authentik.root.test_runner import get_docker_tag
 
-IS_CI = "CI" in environ
-
 
 class DockerTestCase(TestCase):
     """Mixin for dealing with containers"""

tests/e2e/test_endpoints_flow.py

@@ -8,7 +8,8 @@
 from authentik.events.models import Event, EventAction
 from authentik.flows.models import Flow, FlowStageBinding
 from authentik.lib.generators import generate_id
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestEndpointsFlow(SeleniumTestCase):

tests/e2e/test_flows_authenticators.py

@@ -18,7 +18,8 @@
     StaticToken,
 )
 from authentik.stages.authenticator_totp.models import AuthenticatorTOTPStage, TOTPDevice
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestFlowsAuthenticator(SeleniumTestCase):

tests/e2e/test_flows_authenticators_webauthn.py

@@ -15,8 +15,9 @@
     WebAuthnDevice,
 )
 from authentik.stages.identification.models import IdentificationStage
+from tests.decorators import retry
 from tests.e2e.test_flows_login_sfe import login_sfe
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestFlowsAuthenticatorWebAuthn(SeleniumTestCase):

tests/e2e/test_flows_enroll.py

@@ -12,7 +12,8 @@
 from authentik.lib.config import CONFIG
 from authentik.lib.generators import generate_id
 from authentik.stages.identification.models import IdentificationStage
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestFlowsEnroll(SeleniumTestCase):

tests/e2e/test_flows_login.py

@@ -2,7 +2,8 @@
 
 from authentik.blueprints.tests import apply_blueprint
 from authentik.flows.models import Flow
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestFlowsLogin(SeleniumTestCase):

tests/e2e/test_flows_login_sfe.py

@@ -10,7 +10,8 @@
 from authentik.core.models import User
 from authentik.flows.models import NotConfiguredAction
 from authentik.stages.authenticator_validate.models import AuthenticatorValidateStage, DeviceClasses
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 def login_sfe(driver: WebDriver, user: User):

tests/e2e/test_flows_recovery.py

@@ -13,7 +13,8 @@
 from authentik.lib.config import CONFIG
 from authentik.lib.generators import generate_id
 from authentik.stages.identification.models import IdentificationStage
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestFlowsRecovery(SeleniumTestCase):

tests/e2e/test_flows_stage_setup.py

@@ -8,7 +8,8 @@
 from authentik.flows.models import Flow, FlowDesignation
 from authentik.lib.generators import generate_key
 from authentik.stages.password.models import PasswordStage
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.selenium import SeleniumTestCase
 
 
 class TestFlowsStageSetup(SeleniumTestCase):