diff --git a/posthog/api/test/test_preflight.py b/posthog/api/test/test_preflight.py index 6b63ad28f7542..7518104f7cb7e 100644 --- a/posthog/api/test/test_preflight.py +++ b/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 diff --git a/posthog/cloud_utils.py b/posthog/cloud_utils.py index cc0aec67d15ff..151320df32629 100644 --- a/posthog/cloud_utils.py +++ b/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"]: diff --git a/posthog/conftest.py b/posthog/conftest.py index 2bd97894ddf09..d1feb049174e4 100644 --- a/posthog/conftest.py +++ b/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() diff --git a/posthog/session_recordings/queries/test/test_session_recording_list_from_session_replay.py b/posthog/session_recordings/queries/test/test_session_recording_list_from_session_replay.py index b86a77dfe0522..b0020d2a5de48 100644 --- a/posthog/session_recordings/queries/test/test_session_recording_list_from_session_replay.py +++ b/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"): diff --git a/posthog/settings/base_variables.py b/posthog/settings/base_variables.py index fa63b80da1789..9bbc3c583ab82 100644 --- a/posthog/settings/base_variables.py +++ b/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 diff --git a/posthog/test/base.py b/posthog/test/base.py index 255ccc5ef38fd..993a267b6d7d1 100644 --- a/posthog/test/base.py +++ b/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] = "user1@posthog.com" 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): diff --git a/posthog/test/test_cloud_utils.py b/posthog/test/test_cloud_utils.py index a8c6eb5d3c9b9..efa080ff56fec 100644 --- a/posthog/test/test_cloud_utils.py +++ b/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() diff --git a/posthog/test/test_middleware.py b/posthog/test/test_middleware.py index c75f47a4cf69f..d6d72b5d0f698 100644 --- a/posthog/test/test_middleware.py +++ b/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): diff --git a/posthog/utils.py b/posthog/utils.py index 57415aa4ee0ba..72a246e3b0bb5 100644 --- a/posthog/utils.py +++ b/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: diff --git a/posthog/warehouse/api/test/test_external_data_source.py b/posthog/warehouse/api/test/test_external_data_source.py index 07ce645788d58..b115d4df4a520 100644 --- a/posthog/warehouse/api/test/test_external_data_source.py +++ b/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(