core: remove session migration

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

authentik/core/migrations/0046_session_and_more.py

@@ -1,126 +1,9 @@
 # Generated by Django 5.0.11 on 2025-01-27 12:58
 
 import uuid
-import pickle  # nosec
-from django.core import signing
-from django.contrib.auth import BACKEND_SESSION_KEY, HASH_SESSION_KEY, SESSION_KEY
 from django.db import migrations, models
 import django.db.models.deletion
 from django.conf import settings
-from django.contrib.sessions.backends.cache import KEY_PREFIX
-from django.utils.timezone import now, timedelta
-from authentik.lib.migrations import progress_bar
-from authentik.root.middleware import ClientIPMiddleware
-
-
-SESSION_CACHE_ALIAS = "default"
-
-
-class PickleSerializer:
-    """
-    Simple wrapper around pickle to be used in signing.dumps()/loads() and
-    cache backends.
-    """
-
-    def __init__(self, protocol=None):
-        self.protocol = pickle.HIGHEST_PROTOCOL if protocol is None else protocol
-
-    def dumps(self, obj):
-        """Pickle data to be stored in redis"""
-        return pickle.dumps(obj, self.protocol)
-
-    def loads(self, data):
-        """Unpickle data to be loaded from redis"""
-        try:
-            return pickle.loads(data)  # nosec
-        except Exception:
-            return {}
-
-
-def _migrate_session(
-    apps,
-    db_alias,
-    session_key,
-    session_data,
-    expires,
-):
-    Session = apps.get_model("authentik_core", "Session")
-    OldAuthenticatedSession = apps.get_model("authentik_core", "OldAuthenticatedSession")
-    AuthenticatedSession = apps.get_model("authentik_core", "AuthenticatedSession")
-
-    old_auth_session = (
-        OldAuthenticatedSession.objects.using(db_alias).filter(session_key=session_key).first()
-    )
-
-    args = {
-        "session_key": session_key,
-        "expires": expires,
-        "last_ip": ClientIPMiddleware.default_ip,
-        "last_user_agent": "",
-        "session_data": {},
-    }
-    for k, v in session_data.items():
-        if k == "authentik/stages/user_login/last_ip":
-            args["last_ip"] = v
-        elif k in ["last_user_agent", "last_used"]:
-            args[k] = v
-        elif args in [SESSION_KEY, BACKEND_SESSION_KEY, HASH_SESSION_KEY]:
-            pass
-        else:
-            args["session_data"][k] = v
-    if old_auth_session:
-        args["last_user_agent"] = old_auth_session.last_user_agent
-        args["last_used"] = old_auth_session.last_used
-
-    args["session_data"] = pickle.dumps(args["session_data"])
-    session = Session.objects.using(db_alias).create(**args)
-
-    if old_auth_session:
-        AuthenticatedSession.objects.using(db_alias).create(
-            session=session,
-            user=old_auth_session.user,
-        )
-
-
-def migrate_redis_sessions(apps, schema_editor):
-    from django.core.cache import caches
-
-    db_alias = schema_editor.connection.alias
-    cache = caches[SESSION_CACHE_ALIAS]
-
-    # Not a redis cache, skipping
-    if not hasattr(cache, "keys"):
-        return
-
-    print("\nMigrating Redis sessions to database, this might take a couple of minutes...")
-    for key, session_data in progress_bar(cache.get_many(cache.keys(f"{KEY_PREFIX}*")).items()):
-        _migrate_session(
-            apps=apps,
-            db_alias=db_alias,
-            session_key=key.removeprefix(KEY_PREFIX),
-            session_data=session_data,
-            expires=now() + timedelta(seconds=cache.ttl(key)),
-        )
-
-
-def migrate_database_sessions(apps, schema_editor):
-    DjangoSession = apps.get_model("sessions", "Session")
-    db_alias = schema_editor.connection.alias
-
-    print("\nMigration database sessions, this might take a couple of minutes...")
-    for django_session in progress_bar(DjangoSession.objects.using(db_alias).all()):
-        session_data = signing.loads(
-            django_session.session_data,
-            salt="django.contrib.sessions.SessionStore",
-            serializer=PickleSerializer,
-        )
-        _migrate_session(
-            apps=apps,
-            db_alias=db_alias,
-            session_key=django_session.session_key,
-            session_data=session_data,
-            expires=django_session.expire_date,
-        )
 
 
 class Migration(migrations.Migration):
@@ -230,12 +113,4 @@ class Migration(migrations.Migration):
                 "verbose_name_plural": "Authenticated Sessions",
             },
         ),
-        migrations.RunPython(
-            code=migrate_redis_sessions,
-            reverse_code=migrations.RunPython.noop,
-        ),
-        migrations.RunPython(
-            code=migrate_database_sessions,
-            reverse_code=migrations.RunPython.noop,
-        ),
     ]

authentik/root/test_runner.py

@@ -11,7 +11,7 @@ from django.test.runner import DiscoverRunner
 from authentik.lib.config import CONFIG
 from authentik.lib.sentry import sentry_init
 from authentik.root.signals import post_startup, pre_startup, startup
-from tests.docker import get_docker_tag
+from tests.e2e.utils import get_docker_tag
 
 # globally set maxDiff to none to show full assert error
 TestCase.maxDiff = None

authentik/sources/kerberos/models.py

@@ -317,7 +317,7 @@ class KerberosSource(Source):
                 usage="accept", name=name, store=self.get_gssapi_store()
             )
         except gssapi.exceptions.GSSError as exc:
-            LOGGER.warn("GSSAPI credentials failure", exc=exc)
+            LOGGER.warning("GSSAPI credentials failure", exc=exc)
             return None
 
 

tests/__init__.py

@@ -1,12 +0,0 @@
-import socket
-from os import environ
-
-IS_CI = "CI" in environ
-RETRIES = int(environ.get("RETRIES", "3")) if IS_CI else 1
-
-
-def get_local_ip() -> str:
-    """Get the local machine's IP"""
-    hostname = socket.gethostname()
-    ip_addr = socket.gethostbyname(hostname)
-    return ip_addr

tests/decorators.py

@@ -1,48 +0,0 @@
-"""authentik e2e testing utilities"""
-
-from collections.abc import Callable
-from functools import wraps
-
-from django.test.testcases import TransactionTestCase
-from selenium.common.exceptions import NoSuchElementException, TimeoutException, WebDriverException
-from structlog.stdlib import get_logger
-
-from tests import RETRIES
-
-
-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,139 +0,0 @@
-"""Docker testing helpers"""
-
-import os
-from time import sleep
-from typing import TYPE_CHECKING, Any
-from unittest.case import TestCase
-
-from docker import DockerClient, from_env
-from docker.errors import DockerException
-from docker.models.containers import Container
-from docker.models.networks import Network
-
-from authentik.lib.generators import generate_id
-from tests import IS_CI
-
-if TYPE_CHECKING:
-    from authentik.outposts.models import Outpost
-
-
-def get_docker_tag() -> str:
-    """Get docker-tag based off of CI variables"""
-    env_pr_branch = "GITHUB_HEAD_REF"
-    default_branch = "GITHUB_REF"
-    branch_name = os.environ.get(default_branch, "main")
-    if os.environ.get(env_pr_branch, "") != "":
-        branch_name = os.environ[env_pr_branch]
-    branch_name = branch_name.replace("refs/heads/", "").replace("/", "-")
-    return f"gh-{branch_name}"
-
-
-class DockerTestCase(TestCase):
-    """Mixin for dealing with containers"""
-
-    max_healthcheck_attempts = 30
-
-    __client: DockerClient
-    __network: Network
-
-    __label_id = generate_id()
-
-    def setUp(self) -> None:
-        self.__client = from_env()
-        self.__network = self.docker_client.networks.create(
-            name=f"authentik-test-{self.__label_id}"
-        )
-        super().setUp()
-
-    @property
-    def docker_client(self) -> DockerClient:
-        return self.__client
-
-    @property
-    def docker_network(self) -> Network:
-        return self.__network
-
-    @property
-    def docker_labels(self) -> dict:
-        return {"io.goauthentik.test": self.__label_id}
-
-    def get_container_image(self, base: str) -> str:
-        """Try to pull docker image based on git branch, fallback to main if not found."""
-        image = f"{base}:gh-main"
-        if not IS_CI:
-            return image
-        try:
-            branch_image = f"{base}:{get_docker_tag()}"
-            self.docker_client.images.pull(branch_image)
-            return branch_image
-        except DockerException:
-            self.docker_client.images.pull(image)
-        return image
-
-    def run_container(self, **specs: dict[str, Any]) -> Container:
-        if "network_mode" not in specs:
-            specs["network"] = self.__network.name
-        specs["labels"] = self.docker_labels
-        specs["detach"] = True
-        if hasattr(self, "live_server_url"):
-            specs.setdefault("environment", {})
-            specs["environment"]["AUTHENTIK_HOST"] = self.live_server_url
-        container = self.docker_client.containers.run(**specs)
-        container.reload()
-        state = container.attrs.get("State", {})
-        if "Health" not in state:
-            return container
-        self.wait_for_container(container)
-        return container
-
-    def output_container_logs(self, container: Container | None = None):
-        """Output the container logs to our STDOUT"""
-        if IS_CI:
-            image = container.image
-            tags = image.tags[0] if len(image.tags) > 0 else str(image)
-            print(f"::group::Container logs - {tags}")
-        for log in container.logs().decode().split("\n"):
-            print(log)
-        if IS_CI:
-            print("::endgroup::")
-
-    def tearDown(self):
-        containers: list[Container] = self.docker_client.containers.list(
-            filters={"label": ",".join(f"{x}={y}" for x, y in self.docker_labels.items())}
-        )
-        for container in containers:
-            self.output_container_logs(container)
-            try:
-                container.stop()
-            except DockerException:
-                pass
-            try:
-                container.remove(force=True)
-            except DockerException:
-                pass
-        self.__network.remove()
-        super().tearDown()
-
-    def wait_for_container(self, container: Container):
-        """Check that container is health"""
-        attempt = 0
-        while attempt < self.max_healthcheck_attempts:
-            container.reload()
-            status = container.attrs.get("State", {}).get("Health", {}).get("Status")
-            if status == "healthy":
-                return container
-            attempt += 1
-            sleep(0.5)
-        self.failureException("Container failed to start")
-
-    def wait_for_outpost(self, outpost: "Outpost"):
-        # Wait until outpost healthcheck succeeds
-        attempt = 0
-        while attempt < self.max_healthcheck_attempts:
-            if len(outpost.state) > 0:
-                state = outpost.state[0]
-                if state.last_seen:
-                    return
-            attempt += 1
-            sleep(0.5)
-        self.failureException("Outpost failed to become healthy")

tests/e2e/test_flows_authenticators.py

@@ -18,12 +18,10 @@ from authentik.stages.authenticator_static.models import (
     StaticToken,
 )
 from authentik.stages.authenticator_totp.models import AuthenticatorTOTPStage, TOTPDevice
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestFlowsAuthenticator(DockerTestCase, SeleniumTestCase):
+class TestFlowsAuthenticator(SeleniumTestCase):
     """test flow with otp stages"""
 
     @retry()

tests/e2e/test_flows_enroll.py

@@ -11,12 +11,10 @@ from authentik.core.models import User
 from authentik.flows.models import Flow
 from authentik.lib.config import CONFIG
 from authentik.stages.identification.models import IdentificationStage
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestFlowsEnroll(DockerTestCase, SeleniumTestCase):
+class TestFlowsEnroll(SeleniumTestCase):
     """Test Enroll flow"""
 
     @retry()

tests/e2e/test_flows_login.py

@@ -2,12 +2,10 @@
 
 from authentik.blueprints.tests import apply_blueprint
 from authentik.flows.models import Flow
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestFlowsLogin(DockerTestCase, SeleniumTestCase):
+class TestFlowsLogin(SeleniumTestCase):
     """test default login flow"""
 
     def tearDown(self):

tests/e2e/test_flows_login_sfe.py

@@ -6,12 +6,10 @@ from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
 
 from authentik.blueprints.tests import apply_blueprint
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestFlowsLoginSFE(DockerTestCase, SeleniumTestCase):
+class TestFlowsLoginSFE(SeleniumTestCase):
     """test default login flow"""
 
     def login(self):

tests/e2e/test_flows_recovery.py

@@ -13,12 +13,10 @@ from authentik.flows.models import Flow
 from authentik.lib.config import CONFIG
 from authentik.lib.generators import generate_id
 from authentik.stages.identification.models import IdentificationStage
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestFlowsRecovery(DockerTestCase, SeleniumTestCase):
+class TestFlowsRecovery(SeleniumTestCase):
     """Test Recovery flow"""
 
     def initial_stages(self, user: User):

tests/e2e/test_flows_stage_setup.py

@@ -8,12 +8,10 @@ from authentik.core.models import User
 from authentik.flows.models import Flow, FlowDesignation
 from authentik.lib.generators import generate_key
 from authentik.stages.password.models import PasswordStage
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestFlowsStageSetup(DockerTestCase, SeleniumTestCase):
+class TestFlowsStageSetup(SeleniumTestCase):
     """test stage setup flows"""
 
     @retry()

tests/e2e/test_provider_ldap.py

@@ -16,12 +16,10 @@ from authentik.lib.generators import generate_id
 from authentik.outposts.apps import MANAGED_OUTPOST
 from authentik.outposts.models import Outpost, OutpostConfig, OutpostType
 from authentik.providers.ldap.models import APIAccessMode, LDAPProvider
-from tests.decorators import retry
-from tests.docker import DockerTestCase
-from tests.websocket import WebsocketTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderLDAP(DockerTestCase, WebsocketTestCase):
+class TestProviderLDAP(SeleniumTestCase):
     """LDAP and Outpost e2e tests"""
 
     def start_ldap(self, outpost: Outpost):

tests/e2e/test_provider_oauth2_github.py

@@ -18,12 +18,10 @@ from authentik.providers.oauth2.models import (
     RedirectURI,
     RedirectURIMatchingMode,
 )
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderOAuth2Github(DockerTestCase, SeleniumTestCase):
+class TestProviderOAuth2Github(SeleniumTestCase):
     """test OAuth Provider flow"""
 
     def setUp(self):

tests/e2e/test_provider_oauth2_grafana.py

@@ -26,12 +26,10 @@ from authentik.providers.oauth2.models import (
     RedirectURIMatchingMode,
     ScopeMapping,
 )
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderOAuth2OAuth(DockerTestCase, SeleniumTestCase):
+class TestProviderOAuth2OAuth(SeleniumTestCase):
     """test OAuth with OAuth Provider flow"""
 
     def setUp(self):

tests/e2e/test_provider_oidc.py

@@ -26,12 +26,10 @@ from authentik.providers.oauth2.models import (
     RedirectURIMatchingMode,
     ScopeMapping,
 )
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderOAuth2OIDC(DockerTestCase, SeleniumTestCase):
+class TestProviderOAuth2OIDC(SeleniumTestCase):
     """test OAuth with OpenID Provider flow"""
 
     def setUp(self):

tests/e2e/test_provider_oidc_implicit.py

@@ -26,12 +26,10 @@ from authentik.providers.oauth2.models import (
     RedirectURIMatchingMode,
     ScopeMapping,
 )
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderOAuth2OIDCImplicit(DockerTestCase, SeleniumTestCase):
+class TestProviderOAuth2OIDCImplicit(SeleniumTestCase):
     """test OAuth with OpenID Provider flow"""
 
     def setUp(self):

tests/e2e/test_provider_proxy.py

@@ -3,8 +3,11 @@
 from base64 import b64encode
 from dataclasses import asdict
 from json import loads
+from sys import platform
 from time import sleep
+from unittest.case import skip, skipUnless
 
+from channels.testing import ChannelsLiveServerTestCase
 from jwt import decode
 from selenium.webdriver.common.by import By
 
@@ -15,13 +18,10 @@ from authentik.lib.generators import generate_id
 from authentik.outposts.models import DockerServiceConnection, Outpost, OutpostConfig, OutpostType
 from authentik.outposts.tasks import outpost_connection_discovery
 from authentik.providers.proxy.models import ProxyProvider
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
-from tests.websocket import WebsocketTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderProxy(DockerTestCase, SeleniumTestCase):
+class TestProviderProxy(SeleniumTestCase):
     """Proxy and Outpost e2e tests"""
 
     def setUp(self):
@@ -37,41 +37,13 @@ class TestProviderProxy(DockerTestCase, SeleniumTestCase):
         """Start proxy container based on outpost created"""
         self.run_container(
             image=self.get_container_image("ghcr.io/goauthentik/dev-proxy"),
-            ports={"9000": "9000"},
-            environment={"AUTHENTIK_TOKEN": outpost.token.key},
+            ports={
+                "9000": "9000",
+            },
+            environment={
+                "AUTHENTIK_TOKEN": outpost.token.key,
+            },
         )
-        self.wait_for_outpost(outpost)
-
-    def _prepare(self):
-        # set additionalHeaders to test later
-        self.user.attributes["additionalHeaders"] = {"X-Foo": "bar"}
-        self.user.save()
-
-        proxy: ProxyProvider = ProxyProvider.objects.create(
-            name=generate_id(),
-            authorization_flow=Flow.objects.get(
-                slug="default-provider-authorization-implicit-consent"
-            ),
-            invalidation_flow=Flow.objects.get(slug="default-provider-invalidation-flow"),
-            internal_host=f"http://{self.host}",
-            external_host="http://localhost:9000",
-            basic_auth_enabled=True,
-            basic_auth_user_attribute="basic-username",
-            basic_auth_password_attribute="basic-password",  # nosec
-        )
-        # Ensure OAuth2 Params are set
-        proxy.set_oauth_defaults()
-        proxy.save()
-        # we need to create an application to actually access the proxy
-        Application.objects.create(name=generate_id(), slug=generate_id(), provider=proxy)
-        outpost: Outpost = Outpost.objects.create(
-            name=generate_id(),
-            type=OutpostType.PROXY,
-        )
-        outpost.providers.add(proxy)
-        outpost.build_user_permissions(outpost.user)
-
-        self.start_proxy(outpost)
 
     @retry()
     @apply_blueprint(
@@ -89,7 +61,44 @@ class TestProviderProxy(DockerTestCase, SeleniumTestCase):
     @reconcile_app("authentik_crypto")
     def test_proxy_simple(self):
         """Test simple outpost setup with single provider"""
-        self._prepare()
+        # set additionalHeaders to test later
+        self.user.attributes["additionalHeaders"] = {"X-Foo": "bar"}
+        self.user.save()
+
+        proxy: ProxyProvider = ProxyProvider.objects.create(
+            name=generate_id(),
+            authorization_flow=Flow.objects.get(
+                slug="default-provider-authorization-implicit-consent"
+            ),
+            invalidation_flow=Flow.objects.get(slug="default-provider-invalidation-flow"),
+            internal_host=f"http://{self.host}",
+            external_host="http://localhost:9000",
+        )
+        # Ensure OAuth2 Params are set
+        proxy.set_oauth_defaults()
+        proxy.save()
+        # we need to create an application to actually access the proxy
+        Application.objects.create(name=generate_id(), slug=generate_id(), provider=proxy)
+        outpost: Outpost = Outpost.objects.create(
+            name=generate_id(),
+            type=OutpostType.PROXY,
+        )
+        outpost.providers.add(proxy)
+        outpost.build_user_permissions(outpost.user)
+
+        self.start_proxy(outpost)
+
+        # Wait until outpost healthcheck succeeds
+        healthcheck_retries = 0
+        while healthcheck_retries < 50:  # noqa: PLR2004
+            if len(outpost.state) > 0:
+                state = outpost.state[0]
+                if state.last_seen:
+                    break
+            healthcheck_retries += 1
+            sleep(0.5)
+        sleep(5)
+
         self.driver.get("http://localhost:9000/api")
         self.login()
         sleep(1)
@@ -128,13 +137,49 @@ class TestProviderProxy(DockerTestCase, SeleniumTestCase):
     @reconcile_app("authentik_crypto")
     def test_proxy_basic_auth(self):
         """Test simple outpost setup with single provider"""
-        self._prepare()
-        # Setup basic auth
         cred = generate_id()
+        attr = "basic-password"  # nosec
         self.user.attributes["basic-username"] = cred
-        self.user.attributes["basic-password"] = cred
+        self.user.attributes[attr] = cred
         self.user.save()
 
+        proxy: ProxyProvider = ProxyProvider.objects.create(
+            name=generate_id(),
+            authorization_flow=Flow.objects.get(
+                slug="default-provider-authorization-implicit-consent"
+            ),
+            invalidation_flow=Flow.objects.get(slug="default-provider-invalidation-flow"),
+            internal_host=f"http://{self.host}",
+            external_host="http://localhost:9000",
+            basic_auth_enabled=True,
+            basic_auth_user_attribute="basic-username",
+            basic_auth_password_attribute=attr,
+        )
+        # Ensure OAuth2 Params are set
+        proxy.set_oauth_defaults()
+        proxy.save()
+        # we need to create an application to actually access the proxy
+        Application.objects.create(name=generate_id(), slug=generate_id(), provider=proxy)
+        outpost: Outpost = Outpost.objects.create(
+            name=generate_id(),
+            type=OutpostType.PROXY,
+        )
+        outpost.providers.add(proxy)
+        outpost.build_user_permissions(outpost.user)
+
+        self.start_proxy(outpost)
+
+        # Wait until outpost healthcheck succeeds
+        healthcheck_retries = 0
+        while healthcheck_retries < 50:  # noqa: PLR2004
+            if len(outpost.state) > 0:
+                state = outpost.state[0]
+                if state.last_seen:
+                    break
+            healthcheck_retries += 1
+            sleep(0.5)
+        sleep(5)
+
         self.driver.get("http://localhost:9000/api")
         self.login()
         sleep(1)
@@ -142,9 +187,9 @@ class TestProviderProxy(DockerTestCase, SeleniumTestCase):
         full_body_text = self.driver.find_element(By.CSS_SELECTOR, "pre").text
         body = loads(full_body_text)
 
-        self.assertEqual(body.get("headers").get("X-Authentik-Username"), [self.user.username])
+        self.assertEqual(body["headers"]["X-Authentik-Username"], [self.user.username])
         auth_header = b64encode(f"{cred}:{cred}".encode()).decode()
-        self.assertEqual(body.get("headers").get("Authorization"), [f"Basic {auth_header}"])
+        self.assertEqual(body["headers"]["Authorization"], [f"Basic {auth_header}"])
 
         self.driver.get("http://localhost:9000/outpost.goauthentik.io/sign_out")
         sleep(2)
@@ -154,7 +199,10 @@ class TestProviderProxy(DockerTestCase, SeleniumTestCase):
         self.assertIn("You've logged out of", title)
 
 
-class TestProviderProxyConnect(DockerTestCase, WebsocketTestCase):
+# TODO: Fix flaky test
+@skip("Flaky test")
+@skipUnless(platform.startswith("linux"), "requires local docker")
+class TestProviderProxyConnect(ChannelsLiveServerTestCase):
     """Test Proxy connectivity over websockets"""
 
     @retry(exceptions=[AssertionError])
@@ -193,7 +241,14 @@ class TestProviderProxyConnect(DockerTestCase, WebsocketTestCase):
         outpost.build_user_permissions(outpost.user)
 
         # Wait until outpost healthcheck succeeds
-        self.wait_for_outpost(outpost)
+        healthcheck_retries = 0
+        while healthcheck_retries < 50:  # noqa: PLR2004
+            if len(outpost.state) > 0:
+                state = outpost.state[0]
+                if state.last_seen and state.version:
+                    break
+            healthcheck_retries += 1
+            sleep(0.5)
 
         state = outpost.state
         self.assertGreaterEqual(len(state), 1)

tests/e2e/test_provider_proxy_forward.py

@@ -13,12 +13,10 @@ from authentik.flows.models import Flow
 from authentik.lib.generators import generate_id
 from authentik.outposts.models import Outpost, OutpostType
 from authentik.providers.proxy.models import ProxyMode, ProxyProvider
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderProxyForward(DockerTestCase, SeleniumTestCase):
+class TestProviderProxyForward(SeleniumTestCase):
     """Proxy and Outpost e2e tests"""
 
     def setUp(self):
@@ -32,11 +30,14 @@ class TestProviderProxyForward(DockerTestCase, SeleniumTestCase):
         """Start proxy container based on outpost created"""
         self.run_container(
             image=self.get_container_image("ghcr.io/goauthentik/dev-proxy"),
-            ports={"9000": "9000"},
-            environment={"AUTHENTIK_TOKEN": outpost.token.key},
+            ports={
+                "9000": "9000",
+            },
+            environment={
+                "AUTHENTIK_TOKEN": outpost.token.key,
+            },
             name="ak-test-outpost",
         )
-        self.wait_for_outpost(outpost)
 
     @apply_blueprint(
         "default/flow-default-authentication-flow.yaml",
@@ -76,6 +77,17 @@ class TestProviderProxyForward(DockerTestCase, SeleniumTestCase):
 
         self.start_outpost(outpost)
 
+        # Wait until outpost healthcheck succeeds
+        healthcheck_retries = 0
+        while healthcheck_retries < 50:  # noqa: PLR2004
+            if len(outpost.state) > 0:
+                state = outpost.state[0]
+                if state.last_seen:
+                    break
+            healthcheck_retries += 1
+            sleep(0.5)
+        sleep(5)
+
     @retry()
     def test_traefik(self):
         """Test traefik"""

tests/e2e/test_provider_radius.py

@@ -1,6 +1,7 @@
 """Radius e2e tests"""
 
 from dataclasses import asdict
+from time import sleep
 
 from pyrad.client import Client
 from pyrad.dictionary import Dictionary
@@ -8,17 +9,14 @@ from pyrad.packet import AccessAccept, AccessReject, AccessRequest
 
 from authentik.blueprints.tests import apply_blueprint
 from authentik.core.models import Application, User
-from authentik.core.tests.utils import create_test_user
 from authentik.flows.models import Flow
 from authentik.lib.generators import generate_id, generate_key
 from authentik.outposts.models import Outpost, OutpostConfig, OutpostType
 from authentik.providers.radius.models import RadiusProvider
-from tests.decorators import retry
-from tests.docker import DockerTestCase
-from tests.websocket import WebsocketTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderRadius(DockerTestCase, WebsocketTestCase):
+class TestProviderRadius(SeleniumTestCase):
     """Radius Outpost e2e tests"""
 
     def setUp(self):
@@ -30,13 +28,13 @@ class TestProviderRadius(DockerTestCase, WebsocketTestCase):
         self.run_container(
             image=self.get_container_image("ghcr.io/goauthentik/dev-radius"),
             ports={"1812/udp": "1812/udp"},
-            environment={"AUTHENTIK_TOKEN": outpost.token.key},
+            environment={
+                "AUTHENTIK_TOKEN": outpost.token.key,
+            },
         )
-        self.wait_for_outpost(outpost)
 
     def _prepare(self) -> User:
         """prepare user, provider, app and container"""
-        self.user = create_test_user()
         radius: RadiusProvider = RadiusProvider.objects.create(
             name=generate_id(),
             authorization_flow=Flow.objects.get(slug="default-authentication-flow"),
@@ -52,6 +50,17 @@ class TestProviderRadius(DockerTestCase, WebsocketTestCase):
         outpost.providers.add(radius)
 
         self.start_radius(outpost)
+
+        # Wait until outpost healthcheck succeeds
+        healthcheck_retries = 0
+        while healthcheck_retries < 50:  # noqa: PLR2004
+            if len(outpost.state) > 0:
+                state = outpost.state[0]
+                if state.last_seen:
+                    break
+            healthcheck_retries += 1
+            sleep(0.5)
+        sleep(5)
         return outpost
 
     @retry()

tests/e2e/test_provider_saml.py

@@ -14,12 +14,10 @@ from authentik.policies.expression.models import ExpressionPolicy
 from authentik.policies.models import PolicyBinding
 from authentik.providers.saml.models import SAMLBindings, SAMLPropertyMapping, SAMLProvider
 from authentik.sources.saml.processors.constants import SAML_BINDING_POST
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestProviderSAML(DockerTestCase, SeleniumTestCase):
+class TestProviderSAML(SeleniumTestCase):
     """test SAML Provider flow"""
 
     def setup_client(self, provider: SAMLProvider, force_post: bool = False):

tests/e2e/test_source_ldap_samba.py

@@ -11,12 +11,10 @@ from authentik.sources.ldap.models import LDAPSource, LDAPSourcePropertyMapping
 from authentik.sources.ldap.sync.groups import GroupLDAPSynchronizer
 from authentik.sources.ldap.sync.membership import MembershipLDAPSynchronizer
 from authentik.sources.ldap.sync.users import UserLDAPSynchronizer
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestSourceLDAPSamba(DockerTestCase, SeleniumTestCase):
+class TestSourceLDAPSamba(SeleniumTestCase):
     """test LDAP Source"""
 
     def setUp(self):

tests/e2e/test_source_oauth_oauth1.py

@@ -16,9 +16,7 @@ from authentik.sources.oauth.models import OAuthSource
 from authentik.sources.oauth.types.registry import SourceType, registry
 from authentik.sources.oauth.views.callback import OAuthCallback
 from authentik.stages.identification.models import IdentificationStage
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
 class OAuth1Callback(OAuthCallback):
@@ -50,7 +48,7 @@ class OAUth1Type(SourceType):
         }
 
 
-class TestSourceOAuth1(DockerTestCase, SeleniumTestCase):
+class TestSourceOAuth1(SeleniumTestCase):
     """Test OAuth1 Source"""
 
     def setUp(self) -> None:

tests/e2e/test_source_oauth_oauth2.py

@@ -16,12 +16,10 @@ from authentik.flows.models import Flow
 from authentik.lib.generators import generate_id
 from authentik.sources.oauth.models import OAuthSource
 from authentik.stages.identification.models import IdentificationStage
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 
-class TestSourceOAuth2(DockerTestCase, SeleniumTestCase):
+class TestSourceOAuth2(SeleniumTestCase):
     """test OAuth Source flow"""
 
     def setUp(self):

tests/e2e/test_source_saml.py

@@ -16,9 +16,7 @@ from authentik.flows.models import Flow
 from authentik.lib.generators import generate_id
 from authentik.sources.saml.models import SAMLBindingTypes, SAMLSource
 from authentik.stages.identification.models import IdentificationStage
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 IDP_CERT = """-----BEGIN CERTIFICATE-----
 MIIDXTCCAkWgAwIBAgIJALmVVuDWu4NYMA0GCSqGSIb3DQEBCwUAMEUxCzAJBgNV
@@ -72,7 +70,7 @@ Sm75WXsflOxuTn08LbgGc4s=
 -----END PRIVATE KEY-----"""
 
 
-class TestSourceSAML(DockerTestCase, SeleniumTestCase):
+class TestSourceSAML(SeleniumTestCase):
     """test SAML Source flow"""
 
     def setUp(self):

tests/e2e/test_source_scim.py

@@ -8,14 +8,12 @@ from docker.types import Healthcheck
 from authentik.lib.generators import generate_id
 from authentik.lib.utils.http import get_http_session
 from authentik.sources.scim.models import SCIMSource
-from tests.browser import SeleniumTestCase
-from tests.decorators import retry
-from tests.docker import DockerTestCase
+from tests.e2e.utils import SeleniumTestCase, retry
 
 TEST_POLL_MAX = 25
 
 
-class TestSourceSCIM(DockerTestCase, SeleniumTestCase):
+class TestSourceSCIM(SeleniumTestCase):
     """test SCIM Source flow"""
 
     def setUp(self):

tests/e2e/utils.py

@@ -1,18 +1,29 @@
 """authentik e2e testing utilities"""
 
-# This file cannot import anything django or anything that will load django
-
 import json
+import os
+import socket
+from collections.abc import Callable
+from functools import lru_cache, wraps
+from os import environ
 from sys import stderr
 from time import sleep
-from typing import TYPE_CHECKING
+from typing import Any
 from unittest.case import TestCase
 from urllib.parse import urlencode
 
+from django.apps import apps
 from django.contrib.staticfiles.testing import StaticLiveServerTestCase
+from django.db import connection
+from django.db.migrations.loader import MigrationLoader
+from django.test.testcases import TransactionTestCase
 from django.urls import reverse
+from docker import DockerClient, from_env
+from docker.errors import DockerException
+from docker.models.containers import Container
+from docker.models.networks import Network
 from selenium import webdriver
-from selenium.common.exceptions import WebDriverException
+from selenium.common.exceptions import NoSuchElementException, TimeoutException, WebDriverException
 from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
 from selenium.webdriver.remote.command import Command
@@ -22,27 +33,137 @@ from selenium.webdriver.support import expected_conditions as ec
 from selenium.webdriver.support.wait import WebDriverWait
 from structlog.stdlib import get_logger
 
-from tests import IS_CI, RETRIES, get_local_ip
-from tests.websocket import BaseWebsocketTestCase
+from authentik.core.api.users import UserSerializer
+from authentik.core.models import User
+from authentik.core.tests.utils import create_test_admin_user
+from authentik.lib.generators import generate_id
 
-if TYPE_CHECKING:
-    from authentik.core.models import User
+IS_CI = "CI" in environ
+RETRIES = int(environ.get("RETRIES", "3")) if IS_CI else 1
 
 
-class BaseSeleniumTestCase(TestCase):
-    """Mixin which adds helpers for spinning up Selenium"""
+def get_docker_tag() -> str:
+    """Get docker-tag based off of CI variables"""
+    env_pr_branch = "GITHUB_HEAD_REF"
+    default_branch = "GITHUB_REF"
+    branch_name = os.environ.get(default_branch, "main")
+    if os.environ.get(env_pr_branch, "") != "":
+        branch_name = os.environ[env_pr_branch]
+    branch_name = branch_name.replace("refs/heads/", "").replace("/", "-")
+    return f"gh-{branch_name}"
+
+
+def get_local_ip() -> str:
+    """Get the local machine's IP"""
+    hostname = socket.gethostname()
+    ip_addr = socket.gethostbyname(hostname)
+    return ip_addr
+
+
+class DockerTestCase(TestCase):
+    """Mixin for dealing with containers"""
+
+    max_healthcheck_attempts = 30
+
+    __client: DockerClient
+    __network: Network
+
+    __label_id = generate_id()
+
+    def setUp(self) -> None:
+        self.__client = from_env()
+        self.__network = self.docker_client.networks.create(name=f"authentik-test-{generate_id()}")
+
+    @property
+    def docker_client(self) -> DockerClient:
+        return self.__client
+
+    @property
+    def docker_network(self) -> Network:
+        return self.__network
+
+    @property
+    def docker_labels(self) -> dict:
+        return {"io.goauthentik.test": self.__label_id}
+
+    def wait_for_container(self, container: Container):
+        """Check that container is health"""
+        attempt = 0
+        while True:
+            container.reload()
+            status = container.attrs.get("State", {}).get("Health", {}).get("Status")
+            if status == "healthy":
+                return container
+            sleep(1)
+            attempt += 1
+            if attempt >= self.max_healthcheck_attempts:
+                self.failureException("Container failed to start")
+
+    def get_container_image(self, base: str) -> str:
+        """Try to pull docker image based on git branch, fallback to main if not found."""
+        image = f"{base}:gh-main"
+        try:
+            branch_image = f"{base}:{get_docker_tag()}"
+            self.docker_client.images.pull(branch_image)
+            return branch_image
+        except DockerException:
+            self.docker_client.images.pull(image)
+        return image
+
+    def run_container(self, **specs: dict[str, Any]) -> Container:
+        if "network_mode" not in specs:
+            specs["network"] = self.__network.name
+        specs["labels"] = self.docker_labels
+        specs["detach"] = True
+        if hasattr(self, "live_server_url"):
+            specs.setdefault("environment", {})
+            specs["environment"]["AUTHENTIK_HOST"] = self.live_server_url
+        container = self.docker_client.containers.run(**specs)
+        container.reload()
+        state = container.attrs.get("State", {})
+        if "Health" not in state:
+            return container
+        self.wait_for_container(container)
+        return container
+
+    def output_container_logs(self, container: Container | None = None):
+        """Output the container logs to our STDOUT"""
+        if IS_CI:
+            image = container.image
+            tags = image.tags[0] if len(image.tags) > 0 else str(image)
+            print(f"::group::Container logs - {tags}")
+        for log in container.logs().decode().split("\n"):
+            print(log)
+        if IS_CI:
+            print("::endgroup::")
+
+    def tearDown(self):
+        containers: list[Container] = self.docker_client.containers.list(
+            filters={"label": ",".join(f"{x}={y}" for x, y in self.docker_labels.items())}
+        )
+        for container in containers:
+            self.output_container_logs(container)
+            try:
+                container.kill()
+            except DockerException:
+                pass
+            try:
+                container.remove(force=True)
+            except DockerException:
+                pass
+        self.__network.remove()
+
+
+class SeleniumTestCase(DockerTestCase, StaticLiveServerTestCase):
+    """StaticLiveServerTestCase which automatically creates a Webdriver instance"""
 
     host = get_local_ip()
     wait_timeout: int
-    user: "User"
+    user: User
 
     def setUp(self):
         if IS_CI:
             print("::group::authentik Logs", file=stderr)
-        from django.apps import apps
-
-        from authentik.core.tests.utils import create_test_admin_user
-
         apps.get_app_config("authentik_tenants").ready()
         self.wait_timeout = 60
         self.driver = self._get_driver()
@@ -169,10 +290,8 @@ class BaseSeleniumTestCase(TestCase):
         password_stage.find_element(By.CSS_SELECTOR, "input[name=password]").send_keys(Keys.ENTER)
         sleep(1)
 
-    def assert_user(self, expected_user: "User"):
+    def assert_user(self, expected_user: User):
         """Check users/me API and assert it matches expected_user"""
-        from authentik.core.api.users import UserSerializer
-
         self.driver.get(self.url("authentik_api:user-me") + "?format=json")
         user_json = self.driver.find_element(By.CSS_SELECTOR, "pre").text
         user = UserSerializer(data=json.loads(user_json)["user"])
@@ -182,9 +301,46 @@ class BaseSeleniumTestCase(TestCase):
         self.assertEqual(user["email"].value, expected_user.email)
 
 
-class SeleniumTestCase(BaseSeleniumTestCase, StaticLiveServerTestCase):
-    """Test case which spins up a selenium instance and a HTTP-only test server"""
+@lru_cache
+def get_loader():
+    """Thin wrapper to lazily get a Migration Loader, only when it's needed
+    and only once"""
+    return MigrationLoader(connection)
 
 
-class WebsocketSeleniumTestCase(BaseSeleniumTestCase, BaseWebsocketTestCase):
-    """Test case which spins up a selenium instance and a Websocket/HTTP test server"""
+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/integration/test_outpost_docker.py

@@ -19,7 +19,7 @@ from authentik.outposts.models import (
 )
 from authentik.outposts.tasks import outpost_connection_discovery
 from authentik.providers.proxy.models import ProxyProvider
-from tests.docker import DockerTestCase, get_docker_tag
+from tests.e2e.utils import DockerTestCase, get_docker_tag
 
 
 class OutpostDockerTests(DockerTestCase, ChannelsLiveServerTestCase):

tests/integration/test_proxy_docker.py

@@ -19,7 +19,7 @@ from authentik.outposts.models import (
 from authentik.outposts.tasks import outpost_connection_discovery
 from authentik.providers.proxy.controllers.docker import DockerController
 from authentik.providers.proxy.models import ProxyProvider
-from tests.docker import DockerTestCase, get_docker_tag
+from tests.e2e.utils import DockerTestCase, get_docker_tag
 
 
 class TestProxyDocker(DockerTestCase, ChannelsLiveServerTestCase):

tests/websocket.py

@@ -1,52 +0,0 @@
-# This file cannot import anything django or anything that will load django
-from sys import stderr
-
-from channels.testing import ChannelsLiveServerTestCase
-from daphne.testing import DaphneProcess
-from structlog.stdlib import get_logger
-
-from tests import IS_CI, get_local_ip
-
-
-def set_database_connection():
-    from django.conf import settings
-
-    settings.DATABASES["default"]["NAME"] = settings.DATABASES["default"]["TEST"]["NAME"]
-    settings.TEST = True
-
-
-class DatabasePatchDaphneProcess(DaphneProcess):
-    # See https://github.com/django/channels/issues/2048
-    # See https://github.com/django/channels/pull/2033
-
-    def __init__(self, host, get_application, kwargs=None, setup=None, teardown=None):
-        super().__init__(host, get_application, kwargs, setup, teardown)
-        self.setup = set_database_connection
-
-
-class BaseWebsocketTestCase(ChannelsLiveServerTestCase):
-    """Base channels test case"""
-
-    host = get_local_ip()
-    ProtocolServerProcess = DatabasePatchDaphneProcess
-
-
-class WebsocketTestCase(BaseWebsocketTestCase):
-    """Test case to allow testing against a running Websocket/HTTP server"""
-
-    def setUp(self):
-        if IS_CI:
-            print("::group::authentik Logs", file=stderr)
-        from django.apps import apps
-
-        from authentik.core.tests.utils import create_test_admin_user
-
-        apps.get_app_config("authentik_tenants").ready()
-        self.logger = get_logger()
-        self.user = create_test_admin_user()
-        super().setUp()
-
-    def tearDown(self):
-        if IS_CI:
-            print("::endgroup::", file=stderr)
-        super().tearDown()