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

@@ -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"

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

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__)))
 

pytest.ini

@@ -1,4 +1,7 @@
 [pytest]
+env =
+    DEBUG=1
+    TEST=1
 DJANGO_SETTINGS_MODULE = posthog.settings
 addopts = -p no:warnings --reuse-db
 

requirements-dev.in

@@ -26,6 +26,7 @@ black
 isort
 pytest
 pytest-django
+pytest-env
 pytest-mock
 pytest-cov
 pytest-split

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