fix: Use env var instead of License to determine instance type (#20189)

* fix: Use env var instead of License to determine instance type

* Improve code formatting

* Fix

* Remove unused `type: ignore`s

* Fix `TestAutoProjectMiddleware`

posthog/api/test/test_preflight.py

@@ -6,7 +6,6 @@
 from rest_framework import status
 
 from posthog.cloud_utils import (
-    TEST_clear_cloud_cache,
     TEST_clear_instance_license_cache,
 )
 from posthog.models.instance_setting import set_instance_setting
@@ -287,20 +286,14 @@ def test_can_create_org_with_multi_org(self):
             self.assertEqual(response.json()["can_create_org"], True)
 
     @pytest.mark.ee
-    def test_cloud_preflight_based_on_license(self):
-        TEST_clear_cloud_cache()
-        try:
-            from ee.models.license import License, LicenseManager
-        except ImportError:
-            pass
-        else:
-            super(LicenseManager, cast(LicenseManager, License.objects)).create(
-                key="key::123",
-                plan="cloud",
-                valid_until=timezone.datetime(2038, 1, 19, 3, 14, 7),
-            )
-
+    def test_cloud_preflight_based_on_region(self):
+        with self.settings(REGION="US"):
             response = self.client.get("/_preflight/")
             assert response.status_code == status.HTTP_200_OK
             assert response.json()["realm"] == "cloud"
-            assert response.json()["cloud"]
+            assert response.json()["cloud"] is True
+        with self.settings(REGION=None):
+            response = self.client.get("/_preflight/")
+            assert response.status_code == status.HTTP_200_OK
+            assert response.json()["realm"] == "hosted-clickhouse"
+            assert response.json()["cloud"] is False

posthog/cloud_utils.py

@@ -12,38 +12,8 @@
 instance_license_cached: Optional["License"] = None
 
 
-# NOTE: This is cached for the lifetime of the instance but this is not an issue as the value is not expected to change
-def is_cloud():
-    global is_cloud_cached
-
-    if not settings.EE_AVAILABLE:
-        return False
-
-    if isinstance(is_cloud_cached, bool):
-        return is_cloud_cached
-
-    if not is_cloud_cached:
-        try:
-            # NOTE: Important not to import this from ee.models as that will cause a circular import for celery
-            from ee.models.license import License
-
-            # TRICKY - The license table may not exist if a migration is running
-            license = License.objects.first_valid()
-            is_cloud_cached = license.plan == "cloud" if license else False
-        except ProgrammingError:
-            # TRICKY - The license table may not exist if a migration is running
-            pass
-        except Exception as e:
-            print("ERROR: Unable to check license", e)  # noqa: T201
-            capture_exception(e)
-
-    return is_cloud_cached
-
-
-# NOTE: This is purely for testing purposes
-def TEST_clear_cloud_cache(value: Optional[bool] = None):
-    global is_cloud_cached
-    is_cloud_cached = value
+def is_cloud() -> bool:
+    return bool(settings.REGION)
 
 
 def get_cached_instance_license() -> Optional["License"]:

posthog/conftest.py

@@ -5,7 +5,7 @@
 from infi.clickhouse_orm import Database
 
 from posthog.client import sync_execute
-from posthog.test.base import TestMixin, run_clickhouse_statement_in_parallel
+from posthog.test.base import PostHogTestCase, run_clickhouse_statement_in_parallel
 
 
 def create_clickhouse_tables(num_tables: int):
@@ -124,7 +124,7 @@ def django_db_setup(django_db_setup, django_db_keepdb):
 
 @pytest.fixture
 def base_test_mixin_fixture():
-    kls = TestMixin()
+    kls = PostHogTestCase()
     kls.setUp()
     kls.setUpTestData()
 

posthog/session_recordings/queries/test/test_session_recording_list_from_session_replay.py

@@ -7,7 +7,6 @@
 
 from posthog.clickhouse.client import sync_execute
 from posthog.clickhouse.log_entries import TRUNCATE_LOG_ENTRIES_TABLE_SQL
-from posthog.cloud_utils import TEST_clear_cloud_cache
 from posthog.constants import AvailableFeature
 from posthog.models import Person, Cohort, GroupTypeMapping
 from posthog.models.action import Action
@@ -690,7 +689,6 @@ def test_event_filter_has_ttl_applied_too(self):
     def test_ttl_days(self):
         assert ttl_days(self.team) == 21
 
-        TEST_clear_cloud_cache()
         with self.is_cloud(True):
             # Far enough in the future from `days_since_blob_ingestion` but not paid
             with freeze_time("2023-09-01T12:00:01Z"):

posthog/settings/base_variables.py

@@ -13,7 +13,9 @@
 DEBUG = get_from_env("DEBUG", False, type_cast=str_to_bool)
 TEST = "test" in sys.argv or sys.argv[0].endswith("pytest") or get_from_env("TEST", False, type_cast=str_to_bool)  # type: bool
 DEMO = get_from_env("DEMO", False, type_cast=str_to_bool)  # Whether this is a managed demo environment
-REGION = get_from_env("REGION", "US")  # Whether this is a Cloud US or Cloud EU instance
+REGION = get_from_env(  # Whether this is the "US" or "EU" Cloud region (REGION's only set on Cloud)
+    "REGION", optional=True
+)
 SELF_CAPTURE = get_from_env("SELF_CAPTURE", DEBUG and not DEMO, type_cast=str_to_bool)
 E2E_TESTING = get_from_env(
     "E2E_TESTING", False, type_cast=str_to_bool

posthog/test/base.py

@@ -21,7 +21,7 @@
 from django.core.cache import cache
 from django.db import connection, connections
 from django.db.migrations.executor import MigrationExecutor
-from django.test import TestCase, TransactionTestCase, override_settings
+from django.test import SimpleTestCase, TestCase, TransactionTestCase, override_settings
 from django.test.utils import CaptureQueriesContext
 from rest_framework.test import APITestCase as DRFTestCase
 
@@ -30,9 +30,7 @@
 from posthog.clickhouse.client.connection import ch_pool
 from posthog.clickhouse.plugin_log_entries import TRUNCATE_PLUGIN_LOG_ENTRIES_TABLE_SQL
 from posthog.cloud_utils import (
-    TEST_clear_cloud_cache,
     TEST_clear_instance_license_cache,
-    is_cloud,
 )
 from posthog.models import Dashboard, DashboardTile, Insight, Organization, Team, User
 from posthog.models.channel_type.sql import (
@@ -186,7 +184,7 @@ def validation_error_response(
         }
 
 
-class TestMixin:
+class PostHogTestCase(SimpleTestCase):
     CONFIG_ORGANIZATION_NAME: str = "Test"
     CONFIG_EMAIL: Optional[str] = "[email protected]"
     CONFIG_PASSWORD: Optional[str] = "testpassword12345"
@@ -234,17 +232,22 @@ def tearDown(self):
             )
         global persons_ordering_int
         persons_ordering_int = 0
-        super().tearDown()  # type: ignore
+        super().tearDown()
 
     def validate_basic_html(self, html_message, site_url, preheader=None):
         # absolute URLs are used
-        self.assertIn(f"{site_url}/static/posthog-logo.png", html_message)  # type: ignore
+        self.assertIn(f"{site_url}/static/posthog-logo.png", html_message)
 
         # CSS is inlined
-        self.assertIn('style="display: none;', html_message)  # type: ignore
+        self.assertIn('style="display: none;', html_message)
 
         if preheader:
-            self.assertIn(preheader, html_message)  # type: ignore
+            self.assertIn(preheader, html_message)
+
+    @contextmanager
+    def is_cloud(self, value: bool):
+        with self.settings(REGION="US" if value else None):
+            yield value
 
     @contextmanager
     def retry_assertion(self, max_retries=5, delay=0.1) -> Generator[None, None, None]:
@@ -296,24 +299,17 @@ def _callTestMethod(self, method):
         )
 
 
-class BaseTest(TestMixin, ErrorResponsesMixin, TestCase):
+class BaseTest(PostHogTestCase, ErrorResponsesMixin, TestCase):
     """
     Base class for performing Postgres-based backend unit tests on.
     Each class and each test is wrapped inside an atomic block to rollback DB commits after each test.
     Read more: https://docs.djangoproject.com/en/3.1/topics/testing/tools/#testcase
     """
 
-    @contextmanager
-    def is_cloud(self, value: bool):
-        previous_value = is_cloud()
-        try:
-            TEST_clear_cloud_cache(value)
-            yield value
-        finally:
-            TEST_clear_cloud_cache(previous_value)
+    pass
 
 
-class NonAtomicBaseTest(TestMixin, ErrorResponsesMixin, TransactionTestCase):
+class NonAtomicBaseTest(PostHogTestCase, ErrorResponsesMixin, TransactionTestCase):
     """
     Django wraps tests in TestCase inside atomic transactions to speed up the run time. TransactionTestCase is the base
     class for TestCase that doesn't implement this atomic wrapper.
@@ -325,18 +321,15 @@ def setUpClass(cls):
         cls.setUpTestData()
 
 
-class APIBaseTest(TestMixin, ErrorResponsesMixin, DRFTestCase):
+class APIBaseTest(PostHogTestCase, ErrorResponsesMixin, DRFTestCase):
     """
     Functional API tests using Django REST Framework test suite.
     """
 
-    initial_cloud_mode: Optional[bool] = False
-
     def setUp(self):
         super().setUp()
 
         cache.clear()
-        TEST_clear_cloud_cache(self.initial_cloud_mode)
         TEST_clear_instance_license_cache()
 
         # Sets the cloud mode to stabilise things tests, especially num query counts
@@ -357,16 +350,6 @@ def assertFasterThan(self, duration_ms: float):
         with assert_faster_than(duration_ms):
             yield
 
-    @contextmanager
-    def is_cloud(self, value: bool):
-        # Typically the is_cloud setting is controlled by License but we need to be able to override it for tests
-        previous_value = is_cloud()
-        try:
-            TEST_clear_cloud_cache(value)
-            yield value
-        finally:
-            TEST_clear_cloud_cache(previous_value)
-
 
 def stripResponse(response, remove=("action", "label", "persons_urls", "filter")):
     if len(response):

posthog/test/test_cloud_utils.py

@@ -4,10 +4,8 @@
 
 from ee.models.license import License
 from posthog.cloud_utils import (
-    TEST_clear_cloud_cache,
     TEST_clear_instance_license_cache,
     get_cached_instance_license,
-    is_cloud,
 )
 from posthog.test.base import BaseTest
 
@@ -16,29 +14,6 @@ class TestCloudUtils(BaseTest):
     def setUp(self):
         assert License.objects.count() == 0
 
-    @pytest.mark.ee
-    def test_is_cloud_returns_correctly(self):
-        TEST_clear_cloud_cache()
-        assert is_cloud() is False
-
-    @pytest.mark.ee
-    def test_is_cloud_checks_license(self):
-        assert is_cloud() is False
-        License.objects.create(key="key", plan="cloud", valid_until=datetime.now() + timedelta(days=30))
-
-        TEST_clear_cloud_cache()
-        assert is_cloud()
-
-    @pytest.mark.ee
-    def test_is_cloud_caches_result(self):
-        TEST_clear_cloud_cache()
-        assert not is_cloud()
-        assert not is_cloud()
-
-        License.objects.create(key="key", plan="cloud", valid_until=datetime.now() + timedelta(days=30))
-
-        assert not is_cloud()
-
     @pytest.mark.ee
     def test_get_cached_instance_license_returns_correctly(self):
         TEST_clear_instance_license_cache()

posthog/test/test_middleware.py

@@ -311,8 +311,8 @@ def test_project_unchanged_when_accessing_missing_project_by_id(self):
         assert response_users_api.json().get("team", {}).get("id") == self.team.id
 
 
+@override_settings(REGION="US")  # As PostHog Cloud
 class TestPostHogTokenCookieMiddleware(APIBaseTest):
-    initial_cloud_mode = True
     CONFIG_AUTO_LOGIN = False
 
     def test_logged_out_client(self):

posthog/utils.py

@@ -815,11 +815,9 @@ def get_instance_realm() -> str:
 
 def get_instance_region() -> Optional[str]:
     """
-    Returns the region for the current instance. `US` or 'EU'.
+    Returns the region for the current Cloud instance. `US` or 'EU'.
     """
-    if is_cloud():
-        return settings.REGION
-    return None
+    return settings.REGION
 
 
 def get_can_create_org(user: Union["AbstractBaseUser", "AnonymousUser"]) -> bool:

posthog/warehouse/api/test/test_external_data_source.py

@@ -7,7 +7,6 @@
 )
 from django.test import override_settings
 from django.conf import settings
-from posthog.cloud_utils import TEST_clear_cloud_cache
 from posthog.models import Team
 import psycopg
 
@@ -194,8 +193,6 @@ def test_database_schema(self):
     def test_internal_postgres(self, patch_get_postgres_schemas):
         patch_get_postgres_schemas.return_value = ["table_1"]
 
-        TEST_clear_cloud_cache(True)
-
         with override_settings(REGION="US"):
             team_2, _ = Team.objects.get_or_create(id=2, organization=self.team.organization)
             response = self.client.get(