From 25afe3b2e468687cdf3abede46da85a211f2e4d4 Mon Sep 17 00:00:00 2001
From: Kim Gustyr <kim.gustyr@flagsmith.com>
Date: Wed, 24 Apr 2024 13:44:30 +0100
Subject: [PATCH] feat: Add Pytest CI mode to optimise migrations (#3815)

---
 api/.env-ci     |  2 +-
 api/.env-local  |  1 -
 api/Makefile    |  2 ++
 api/conftest.py | 50 +++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/api/.env-ci b/api/.env-ci
index ca10efea922f..093c8fb9b364 100644
--- a/api/.env-ci
+++ b/api/.env-ci
@@ -1,5 +1,5 @@
 DATABASE_URL=postgresql://postgres:postgres@localhost:5432/postgres
 ANALYTICS_DATABASE_URL=postgres://postgres:postgres@localhost:5432/analytics
-PYTEST_ADDOPTS=--cov . --cov-report xml -n auto --dist worksteal
+PYTEST_ADDOPTS=--cov . --cov-report xml -n auto --dist worksteal --ci
 GUNICORN_LOGGER_CLASS=util.logging.GunicornJsonCapableLogger
 COVERAGE_CORE=sysmon
diff --git a/api/.env-local b/api/.env-local
index 84b43f8720a3..ff290fe34c85 100644
--- a/api/.env-local
+++ b/api/.env-local
@@ -1,4 +1,3 @@
 DATABASE_URL=postgresql://postgres:password@localhost:5432/flagsmith
 DJANGO_SETTINGS_MODULE=app.settings.local
 PYTEST_ADDOPTS=--cov . --cov-report html -n auto
-GUNICORN_LOGGER_CLASS=util.logging.GunicornJsonCapableLogger
diff --git a/api/Makefile b/api/Makefile
index 26d30475f35c..9bcc5ec18588 100644
--- a/api/Makefile
+++ b/api/Makefile
@@ -9,6 +9,8 @@ DOTENV_OVERRIDE_FILE ?= .env
 
 POETRY_VERSION ?= 1.8.2
 
+GUNICORN_LOGGER_CLASS ?= util.logging.GunicornJsonCapableLogger
+
 -include .env-local
 -include $(DOTENV_OVERRIDE_FILE)
 
diff --git a/api/conftest.py b/api/conftest.py
index 7dd092896230..781171aa148c 100644
--- a/api/conftest.py
+++ b/api/conftest.py
@@ -5,11 +5,15 @@
 import pytest
 from django.contrib.contenttypes.models import ContentType
 from django.core.cache import caches
+from django.db.backends.base.creation import TEST_DATABASE_PREFIX
+from django.test.utils import setup_databases
 from flag_engine.segments.constants import EQUAL
 from moto import mock_dynamodb
 from mypy_boto3_dynamodb.service_resource import DynamoDBServiceResource, Table
+from pytest_django.plugin import blocking_manager_key
 from rest_framework.authtoken.models import Token
 from rest_framework.test import APIClient
+from xdist import get_xdist_worker_id
 
 from api_keys.models import MasterAPIKey
 from environments.identities.models import Identity
@@ -65,6 +69,52 @@
 )
 from users.models import FFAdminUser, UserPermissionGroup
 
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--ci",
+        action="store_true",
+        default=False,
+        help="Enable CI mode",
+    )
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_configure(config: pytest.Config) -> None:
+    if (
+        config.option.ci
+        and config.option.dist != "no"
+        and not hasattr(config, "workerinput")
+    ):
+        with config.stash[blocking_manager_key].unblock():
+            setup_databases(
+                verbosity=config.option.verbose,
+                interactive=False,
+                parallel=config.option.numprocesses,
+            )
+
+
+@pytest.fixture(scope="session")
+def django_db_setup(request: pytest.FixtureRequest) -> None:
+    if (
+        request.config.option.ci
+        # xdist worker id is either `gw[0-9]+` or `master`
+        and (xdist_worker_id_suffix := get_xdist_worker_id(request)[2:]).isnumeric()
+    ):
+        # Django's test database clone indices start at 1,
+        # Pytest's worker indices are 0-based
+        test_db_suffix = str(int(xdist_worker_id_suffix) + 1)
+    else:
+        # Tests are run on main node, which assumes -n0
+        return request.getfixturevalue("django_db_setup")  # pragma: no cover
+
+    from django.conf import settings
+
+    for db_settings in settings.DATABASES.values():
+        test_db_name = f'{TEST_DATABASE_PREFIX}{db_settings["NAME"]}_{test_db_suffix}'
+        db_settings["NAME"] = test_db_name
+
+
 trait_key = "key1"
 trait_value = "value1"