tests: better ws support

cherry-picked from https://github.com/goauthentik/authentik/pull/14539

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

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.e2e.utils import get_docker_tag
+from tests.docker import get_docker_tag
 
 # globally set maxDiff to none to show full assert error
 TestCase.maxDiff = None

tests/__init__.py

@@ -0,0 +1,12 @@
+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/browser.py

@@ -1,29 +1,18 @@
 """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 Any
+from typing import TYPE_CHECKING
 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 NoSuchElementException, TimeoutException, WebDriverException
+from selenium.common.exceptions import WebDriverException
 from selenium.webdriver.common.by import By
 from selenium.webdriver.common.keys import Keys
 from selenium.webdriver.remote.command import Command
@@ -33,137 +22,27 @@ from selenium.webdriver.support import expected_conditions as ec
 from selenium.webdriver.support.wait import WebDriverWait
 from structlog.stdlib import get_logger
 
-from authentik.core.api.users import UserSerializer
+from tests import IS_CI, RETRIES, get_local_ip
-from authentik.core.models import User
+from tests.websocket import BaseWebsocketTestCase
-from authentik.core.tests.utils import create_test_admin_user
-from authentik.lib.generators import generate_id
 
-IS_CI = "CI" in environ
+if TYPE_CHECKING:
-RETRIES = int(environ.get("RETRIES", "3")) if IS_CI else 1
+    from authentik.core.models import User
 
 
-def get_docker_tag() -> str:
+class BaseSeleniumTestCase(TestCase):
-    """Get docker-tag based off of CI variables"""
+    """Mixin which adds helpers for spinning up Selenium"""
-    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()
@@ -290,8 +169,10 @@ class SeleniumTestCase(DockerTestCase, StaticLiveServerTestCase):
         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"])
@@ -301,46 +182,9 @@ class SeleniumTestCase(DockerTestCase, StaticLiveServerTestCase):
         self.assertEqual(user["email"].value, expected_user.email)
 
 
-@lru_cache
+class SeleniumTestCase(BaseSeleniumTestCase, StaticLiveServerTestCase):
-def get_loader():
+    """Test case which spins up a selenium instance and a HTTP-only test server"""
-    """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):
+class WebsocketSeleniumTestCase(BaseSeleniumTestCase, BaseWebsocketTestCase):
-    """Retry test multiple times. Default to catching Selenium Timeout Exception"""
+    """Test case which spins up a selenium instance and a Websocket/HTTP test server"""
-
-    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/decorators.py

@@ -0,0 +1,48 @@
+"""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

@@ -0,0 +1,139 @@
+"""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,10 +18,12 @@ from authentik.stages.authenticator_static.models import (
     StaticToken,
 )
 from authentik.stages.authenticator_totp.models import AuthenticatorTOTPStage, TOTPDevice
-from tests.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestFlowsAuthenticator(SeleniumTestCase):
+class TestFlowsAuthenticator(DockerTestCase, SeleniumTestCase):
     """test flow with otp stages"""
 
     @retry()

tests/e2e/test_flows_enroll.py

@@ -11,10 +11,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestFlowsEnroll(SeleniumTestCase):
+class TestFlowsEnroll(DockerTestCase, SeleniumTestCase):
     """Test Enroll flow"""
 
     @retry()

tests/e2e/test_flows_login.py

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

tests/e2e/test_flows_login_sfe.py

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

tests/e2e/test_flows_recovery.py

@@ -13,10 +13,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestFlowsRecovery(SeleniumTestCase):
+class TestFlowsRecovery(DockerTestCase, SeleniumTestCase):
     """Test Recovery flow"""
 
     def initial_stages(self, user: User):

tests/e2e/test_flows_stage_setup.py

@@ -8,10 +8,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestFlowsStageSetup(SeleniumTestCase):
+class TestFlowsStageSetup(DockerTestCase, SeleniumTestCase):
     """test stage setup flows"""
 
     @retry()

tests/e2e/test_provider_ldap.py

@@ -16,10 +16,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.docker import DockerTestCase
+from tests.websocket import WebsocketTestCase
 
 
-class TestProviderLDAP(SeleniumTestCase):
+class TestProviderLDAP(DockerTestCase, WebsocketTestCase):
     """LDAP and Outpost e2e tests"""
 
     def start_ldap(self, outpost: Outpost):

tests/e2e/test_provider_oauth2_github.py

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

tests/e2e/test_provider_oauth2_grafana.py

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

tests/e2e/test_provider_oidc.py

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

tests/e2e/test_provider_oidc_implicit.py

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

tests/e2e/test_provider_proxy.py

@@ -3,11 +3,8 @@
 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
 
@@ -18,10 +15,13 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
+from tests.websocket import WebsocketTestCase
 
 
-class TestProviderProxy(SeleniumTestCase):
+class TestProviderProxy(DockerTestCase, SeleniumTestCase):
     """Proxy and Outpost e2e tests"""
 
     def setUp(self):
@@ -37,13 +37,41 @@ class TestProviderProxy(SeleniumTestCase):
         """Start proxy container based on outpost created"""
         self.run_container(
             image=self.get_container_image("ghcr.io/goauthentik/dev-proxy"),
-            ports={
+            ports={"9000": "9000"},
-                "9000": "9000",
+            environment={"AUTHENTIK_TOKEN": outpost.token.key},
-            },
-            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(
@@ -61,44 +89,7 @@ class TestProviderProxy(SeleniumTestCase):
     @reconcile_app("authentik_crypto")
     def test_proxy_simple(self):
         """Test simple outpost setup with single provider"""
-        # set additionalHeaders to test later
+        self._prepare()
-        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)
@@ -137,49 +128,13 @@ class TestProviderProxy(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[attr] = cred
+        self.user.attributes["basic-password"] = 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)
@@ -187,9 +142,9 @@ class TestProviderProxy(SeleniumTestCase):
         full_body_text = self.driver.find_element(By.CSS_SELECTOR, "pre").text
         body = loads(full_body_text)
 
-        self.assertEqual(body["headers"]["X-Authentik-Username"], [self.user.username])
+        self.assertEqual(body.get("headers").get("X-Authentik-Username"), [self.user.username])
         auth_header = b64encode(f"{cred}:{cred}".encode()).decode()
-        self.assertEqual(body["headers"]["Authorization"], [f"Basic {auth_header}"])
+        self.assertEqual(body.get("headers").get("Authorization"), [f"Basic {auth_header}"])
 
         self.driver.get("http://localhost:9000/outpost.goauthentik.io/sign_out")
         sleep(2)
@@ -199,10 +154,7 @@ class TestProviderProxy(SeleniumTestCase):
         self.assertIn("You've logged out of", title)
 
 
-# TODO: Fix flaky test
+class TestProviderProxyConnect(DockerTestCase, WebsocketTestCase):
-@skip("Flaky test")
-@skipUnless(platform.startswith("linux"), "requires local docker")
-class TestProviderProxyConnect(ChannelsLiveServerTestCase):
     """Test Proxy connectivity over websockets"""
 
     @retry(exceptions=[AssertionError])
@@ -241,14 +193,7 @@ class TestProviderProxyConnect(ChannelsLiveServerTestCase):
         outpost.build_user_permissions(outpost.user)
 
         # Wait until outpost healthcheck succeeds
-        healthcheck_retries = 0
+        self.wait_for_outpost(outpost)
-        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,10 +13,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestProviderProxyForward(SeleniumTestCase):
+class TestProviderProxyForward(DockerTestCase, SeleniumTestCase):
     """Proxy and Outpost e2e tests"""
 
     def setUp(self):
@@ -30,14 +32,11 @@ class TestProviderProxyForward(SeleniumTestCase):
         """Start proxy container based on outpost created"""
         self.run_container(
             image=self.get_container_image("ghcr.io/goauthentik/dev-proxy"),
-            ports={
+            ports={"9000": "9000"},
-                "9000": "9000",
+            environment={"AUTHENTIK_TOKEN": outpost.token.key},
-            },
-            environment={
-                "AUTHENTIK_TOKEN": outpost.token.key,
-            },
             name="ak-test-outpost",
         )
+        self.wait_for_outpost(outpost)
 
     @apply_blueprint(
         "default/flow-default-authentication-flow.yaml",
@@ -77,17 +76,6 @@ class TestProviderProxyForward(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,7 +1,6 @@
 """Radius e2e tests"""
 
 from dataclasses import asdict
-from time import sleep
 
 from pyrad.client import Client
 from pyrad.dictionary import Dictionary
@@ -9,14 +8,17 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.decorators import retry
+from tests.docker import DockerTestCase
+from tests.websocket import WebsocketTestCase
 
 
-class TestProviderRadius(SeleniumTestCase):
+class TestProviderRadius(DockerTestCase, WebsocketTestCase):
     """Radius Outpost e2e tests"""
 
     def setUp(self):
@@ -28,13 +30,13 @@ class TestProviderRadius(SeleniumTestCase):
         self.run_container(
             image=self.get_container_image("ghcr.io/goauthentik/dev-radius"),
             ports={"1812/udp": "1812/udp"},
-            environment={
+            environment={"AUTHENTIK_TOKEN": outpost.token.key},
-                "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"),
@@ -50,17 +52,6 @@ class TestProviderRadius(SeleniumTestCase):
         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,10 +14,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestProviderSAML(SeleniumTestCase):
+class TestProviderSAML(DockerTestCase, SeleniumTestCase):
     """test SAML Provider flow"""
 
     def setup_client(self, provider: SAMLProvider, force_post: bool = False):

tests/e2e/test_source_ldap_samba.py

@@ -11,10 +11,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestSourceLDAPSamba(SeleniumTestCase):
+class TestSourceLDAPSamba(DockerTestCase, SeleniumTestCase):
     """test LDAP Source"""
 
     def setUp(self):

tests/e2e/test_source_oauth_oauth1.py

@@ -16,7 +16,9 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
 class OAuth1Callback(OAuthCallback):
@@ -48,7 +50,7 @@ class OAUth1Type(SourceType):
         }
 
 
-class TestSourceOAuth1(SeleniumTestCase):
+class TestSourceOAuth1(DockerTestCase, SeleniumTestCase):
     """Test OAuth1 Source"""
 
     def setUp(self) -> None:

tests/e2e/test_source_oauth_oauth2.py

@@ -16,10 +16,12 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 
-class TestSourceOAuth2(SeleniumTestCase):
+class TestSourceOAuth2(DockerTestCase, SeleniumTestCase):
     """test OAuth Source flow"""
 
     def setUp(self):

tests/e2e/test_source_saml.py

@@ -16,7 +16,9 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 IDP_CERT = """-----BEGIN CERTIFICATE-----
 MIIDXTCCAkWgAwIBAgIJALmVVuDWu4NYMA0GCSqGSIb3DQEBCwUAMEUxCzAJBgNV
@@ -70,7 +72,7 @@ Sm75WXsflOxuTn08LbgGc4s=
 -----END PRIVATE KEY-----"""
 
 
-class TestSourceSAML(SeleniumTestCase):
+class TestSourceSAML(DockerTestCase, SeleniumTestCase):
     """test SAML Source flow"""
 
     def setUp(self):

tests/e2e/test_source_scim.py

@@ -8,12 +8,14 @@ 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.e2e.utils import SeleniumTestCase, retry
+from tests.browser import SeleniumTestCase
+from tests.decorators import retry
+from tests.docker import DockerTestCase
 
 TEST_POLL_MAX = 25
 
 
-class TestSourceSCIM(SeleniumTestCase):
+class TestSourceSCIM(DockerTestCase, SeleniumTestCase):
     """test SCIM Source flow"""
 
     def setUp(self):

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.e2e.utils import DockerTestCase, get_docker_tag
+from tests.docker 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.e2e.utils import DockerTestCase, get_docker_tag
+from tests.docker import DockerTestCase, get_docker_tag
 
 
 class TestProxyDocker(DockerTestCase, ChannelsLiveServerTestCase):

tests/websocket.py

@@ -0,0 +1,52 @@
+# 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()