From b865e9797315d0177a1bbfefc6be5d10212ee3ea Mon Sep 17 00:00:00 2001
From: Marc 'risson' Schmitt <marc.schmitt@risson.space>
Date: Sun, 16 Feb 2025 17:35:38 +0100
Subject: [PATCH] ci: parallelize unit tests (#13036)

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>
---
 .github/workflows/ci-main.yml | 33 ++++++++++++++++++++++++++-------
 Makefile                      |  5 +++++
 authentik/root/test_plugin.py | 18 ++++++++++++++++++
 3 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/ci-main.yml b/.github/workflows/ci-main.yml
index d88da9dd6f..982abbff1d 100644
--- a/.github/workflows/ci-main.yml
+++ b/.github/workflows/ci-main.yml
@@ -43,15 +43,26 @@ jobs:
         uses: ./.github/actions/setup
       - name: run migrations
         run: poetry run python -m lifecycle.migrate
-  test-migrations-from-stable:
-    name: test-migrations-from-stable - PostgreSQL ${{ matrix.psql }}
+  test-make-seed:
     runs-on: ubuntu-latest
+    steps:
+      - id: seed
+        run: |
+          echo "seed=$(printf "%d\n" "0x$(openssl rand -hex 4)")" >> "$GITHUB_OUTPUT"
+    outputs:
+      seed: ${{ steps.seed.outputs.seed }}
+  test-migrations-from-stable:
+    name: test-migrations-from-stable - PostgreSQL ${{ matrix.psql }} - Run ${{ matrix.run_id }}/5
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    needs: test-make-seed
     strategy:
       fail-fast: false
       matrix:
         psql:
           - 15-alpine
           - 16-alpine
+        run_id: [1, 2, 3, 4, 5]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -93,18 +104,23 @@ jobs:
         env:
           # Test in the main database that we just migrated from the previous stable version
           AUTHENTIK_POSTGRESQL__TEST__NAME: authentik
+          CI_TEST_SEED: ${{ needs.test-make-seed.outputs.seed }}
+          CI_RUN_ID: ${{ matrix.run_id }}
+          CI_TOTAL_RUNS: "5"
         run: |
-          poetry run make test
+          poetry run make ci-test
   test-unittest:
-    name: test-unittest - PostgreSQL ${{ matrix.psql }}
+    name: test-unittest - PostgreSQL ${{ matrix.psql }} - Run ${{ matrix.run_id }}/5
     runs-on: ubuntu-latest
-    timeout-minutes: 30
+    timeout-minutes: 20
+    needs: test-make-seed
     strategy:
       fail-fast: false
       matrix:
         psql:
           - 15-alpine
           - 16-alpine
+        run_id: [1, 2, 3, 4, 5]
     steps:
       - uses: actions/checkout@v4
       - name: Setup authentik env
@@ -112,9 +128,12 @@ jobs:
         with:
           postgresql_version: ${{ matrix.psql }}
       - name: run unittest
+        env:
+          CI_TEST_SEED: ${{ needs.test-make-seed.outputs.seed }}
+          CI_RUN_ID: ${{ matrix.run_id }}
+          CI_TOTAL_RUNS: "5"
         run: |
-          poetry run make test
-          poetry run coverage xml
+          poetry run make ci-test
       - if: ${{ always() }}
         uses: codecov/codecov-action@v5
         with:
diff --git a/Makefile b/Makefile
index 29dccf8a89..c504d228d7 100644
--- a/Makefile
+++ b/Makefile
@@ -284,3 +284,8 @@ ci-bandit: ci--meta-debug
 
 ci-pending-migrations: ci--meta-debug
 	ak makemigrations --check
+
+ci-test: ci--meta-debug
+	coverage run manage.py test --keepdb --randomly-seed ${CI_TEST_SEED} authentik
+	coverage report
+	coverage xml
diff --git a/authentik/root/test_plugin.py b/authentik/root/test_plugin.py
index 815b8e485b..2fed148e08 100644
--- a/authentik/root/test_plugin.py
+++ b/authentik/root/test_plugin.py
@@ -1,3 +1,4 @@
+import math
 from os import environ
 from ssl import OPENSSL_VERSION
 
@@ -24,3 +25,20 @@ def pytest_report_header(*_, **__):
         f"authentik version: {get_full_version()}",
         f"OpenSSL version: {OPENSSL_VERSION}, FIPS: {backend._fips_enabled}",
     ]
+
+
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
+    current_id = int(environ.get("CI_RUN_ID", 0)) - 1
+    total_ids = int(environ.get("CI_TOTAL_RUNS", 0))
+
+    if total_ids:
+        num_tests = len(items)
+        matrix_size = math.ceil(num_tests / total_ids)
+
+        start = current_id * matrix_size
+        end = (current_id + 1) * matrix_size
+
+        deselected_items = items[:start] + items[end:]
+        config.hook.pytest_deselected(items=deselected_items)
+        items[:] = items[start:end]
+        print(f" Executing {start} - {end} tests")