From 50b12d9bb7629fbafc55f73308558513fd27a407 Mon Sep 17 00:00:00 2001
From: Paul D'Ambra <paul@posthog.com>
Date: Sat, 20 Nov 2021 18:18:39 +0100
Subject: [PATCH] Set debug, test, and primary_db env variables based on
 context (#7250)

* set debug, test, and primary_db env variables based on context

* remove unnecessary print

* also use pytest.ini for some tests run straight from pytest at the CLI
---
 bin/tests            |  2 --
 ee/pytest.ini        | 12 ++++++++++++
 posthog/settings.py  | 43 +++++++++++++++++++++++++++++++++++++++++++
 pytest.ini           |  3 +++
 requirements-dev.in  |  1 +
 requirements-dev.txt |  3 +++
 6 files changed, 62 insertions(+), 2 deletions(-)
 create mode 100644 ee/pytest.ini

diff --git a/bin/tests b/bin/tests
index e5194fb8524..36bcf34eb0b 100755
--- a/bin/tests
+++ b/bin/tests
@@ -23,7 +23,5 @@ if [ $# -eq 0 ]; then
 fi
 
 export REDIS_URL='redis:///'
-export TEST=1
-export DEBUG=1
 psql posthog -c "drop database if exists test_posthog"
 nodemon -w ./posthog -w ./ee --ext py --exec "OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES pytest --reuse-db -s $* --snapshot-update; mypy posthog ee"
diff --git a/ee/pytest.ini b/ee/pytest.ini
new file mode 100644
index 00000000000..b8fa0695ddf
--- /dev/null
+++ b/ee/pytest.ini
@@ -0,0 +1,12 @@
+[pytest]
+env =
+    PRIMARY_DB=clickhouse
+    DEBUG=1
+    TEST=1
+DJANGO_SETTINGS_MODULE = posthog.settings
+addopts = -p no:warnings --reuse-db
+
+markers =
+    ee
+    clickhouse_only
+    skip_on_multitenancy
diff --git a/posthog/settings.py b/posthog/settings.py
index 08f5fef6168..0c486db4e11 100644
--- a/posthog/settings.py
+++ b/posthog/settings.py
@@ -51,6 +51,49 @@ def get_list(text: str) -> List[str]:
     return [item.strip() for item in text.split(",")]
 
 
+"""
+There are several options:
+1) running in pycharm
+        second argument is "test"
+2) running pytest at the CLI
+        first argument is the path to pytest and ends pytest
+3) running pytest using the script at /bin/tests
+        first argument is the path to pytest and ends pytest
+4) running in some other context (e.g. in prod)
+        first argument does not end pytest
+        second argument is not test
+
+Arguments to the application will be slightly different in each case
+
+So, in order to set test variables we need to look in slightly different places 
+
+The /bin/tests file also runs mypy to do type checking. This needs DEBUG=1 set too
+
+Running pytest directly does not always load django settings but sometimes needs these environment variables.
+We use pytest-env to let us set environment variables from the closest pytest.ini 
+
+We can't rely only on pytest.ini as some tests evaluate this file before its environment variables have been read
+"""
+runner = sys.argv[0] if len(sys.argv) >= 1 else None
+
+if runner:
+    cmd = sys.argv[1] if len(sys.argv) >= 2 else None
+
+    if cmd == "test" or runner.endswith("pytest") or runner.endswith("mypy"):
+        print("Running in test mode. Setting DEBUG and TEST environment variables.")
+        os.environ["DEBUG"] = "1"
+        os.environ["TEST"] = "1"
+
+        try:
+            path = sys.argv[2] if cmd == "test" else sys.argv[3]
+            if path.startswith("ee"):
+                print("Running EE tests. Setting clickhouse as primary database.")
+                os.environ["PRIMARY_DB"] = "clickhouse"
+        except IndexError:
+            # there was no path, we don't want to set PRIMARY_DB
+            pass
+
+
 # Build paths inside the project like this: os.path.join(BASE_DIR, ...)
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
diff --git a/pytest.ini b/pytest.ini
index 3dbc6d16589..4af882084e4 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -1,4 +1,7 @@
 [pytest]
+env =
+    DEBUG=1
+    TEST=1
 DJANGO_SETTINGS_MODULE = posthog.settings
 addopts = -p no:warnings --reuse-db
 
diff --git a/requirements-dev.in b/requirements-dev.in
index 9cfd2039147..4c29cedfd8a 100644
--- a/requirements-dev.in
+++ b/requirements-dev.in
@@ -26,6 +26,7 @@ black
 isort
 pytest
 pytest-django
+pytest-env
 pytest-mock
 pytest-cov
 pytest-split
diff --git a/requirements-dev.txt b/requirements-dev.txt
index caf88549ab2..13d92b3f4eb 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -127,6 +127,8 @@ pytest-cov==2.12.1
     # via -r requirements-dev.in
 pytest-django==4.1.0
     # via -r requirements-dev.in
+pytest-env==0.6.2
+    # via -r requirements-dev.in
 pytest-mock==3.5.1
     # via -r requirements-dev.in
 pytest-split==0.3.3
@@ -136,6 +138,7 @@ pytest==6.2.2
     #   -r requirements-dev.in
     #   pytest-cov
     #   pytest-django
+    #   pytest-env
     #   pytest-mock
     #   pytest-split
     #   syrupy