From 8fe7a078c30ce89769aa3af03f1e045e3f1f4b91 Mon Sep 17 00:00:00 2001 From: Michael Matloka Date: Tue, 10 Sep 2024 17:32:16 +0200 Subject: [PATCH] Restore `is_demo` support and test --- ee/api/test/test_team.py | 6 ++-- posthog/api/project.py | 5 ++- posthog/api/signup.py | 2 +- posthog/api/team.py | 6 +++- posthog/api/test/test_routing.py | 4 ++- posthog/api/test/test_team.py | 19 +++++++--- posthog/demo/legacy/__init__.py | 15 -------- posthog/models/organization.py | 4 ++- posthog/models/project.py | 14 +++++--- posthog/models/team/team.py | 56 +++++++++++++++++++---------- posthog/models/test/test_project.py | 4 +++ posthog/models/user.py | 4 ++- posthog/test/base.py | 1 + posthog/test/test_team.py | 18 +++++----- 14 files changed, 98 insertions(+), 60 deletions(-) diff --git a/ee/api/test/test_team.py b/ee/api/test/test_team.py index 254b6aa337382d..d90f699fcd5c69 100644 --- a/ee/api/test/test_team.py +++ b/ee/api/test/test_team.py @@ -56,7 +56,7 @@ def test_create_demo_team(self, *args): self.organization_membership.level = OrganizationMembership.Level.ADMIN self.organization_membership.save() response = self.client.post("/api/environments/", {"name": "Hedgebox", "is_demo": True}) - self.assertEqual(Team.objects.count(), 2) + self.assertEqual(Team.objects.count(), 3) self.assertEqual(response.status_code, 201) response_data = response.json() self.assertDictContainsSubset( @@ -73,7 +73,7 @@ def test_create_two_demo_teams(self, *args): self.organization_membership.level = OrganizationMembership.Level.ADMIN self.organization_membership.save() response = self.client.post("/api/environments/", {"name": "Hedgebox", "is_demo": True}) - self.assertEqual(Team.objects.count(), 2) + self.assertEqual(Team.objects.count(), 3) self.assertEqual(response.status_code, 201) response_data = response.json() self.assertDictContainsSubset( @@ -85,7 +85,7 @@ def test_create_two_demo_teams(self, *args): response_data, ) response_2 = self.client.post("/api/environments/", {"name": "Hedgebox", "is_demo": True}) - self.assertEqual(Team.objects.count(), 2) + self.assertEqual(Team.objects.count(), 3) response_2_data = response_2.json() self.assertDictContainsSubset( { diff --git a/posthog/api/project.py b/posthog/api/project.py index ca61df684f8012..ef1c54faae677d 100644 --- a/posthog/api/project.py +++ b/posthog/api/project.py @@ -210,7 +210,10 @@ def create(self, validated_data: dict[str, Any], **kwargs) -> Project: if field_name in self.Meta.team_passthrough_fields: team_fields[field_name] = validated_data.pop(field_name) project, team = Project.objects.create_with_team( - organization_id=self.context["view"].organization_id, **validated_data, team_fields=team_fields + organization_id=self.context["view"].organization_id, + initiating_user=self.context["request"].user, + **validated_data, + team_fields=team_fields, ) request.user.current_team = team diff --git a/posthog/api/signup.py b/posthog/api/signup.py index 7cda79d66195de..3847999ec551d1 100644 --- a/posthog/api/signup.py +++ b/posthog/api/signup.py @@ -161,7 +161,7 @@ def enter_demo(self, validated_data) -> User: return self._user def create_team(self, organization: Organization, user: User) -> Team: - return Team.objects.create_with_data(user=user, organization=organization) + return Team.objects.create_with_data(initiating_user=user, organization=organization) def to_representation(self, instance) -> dict: data = UserBasicSerializer(instance=instance).data diff --git a/posthog/api/team.py b/posthog/api/team.py index 4711ef6addb4d9..34349958a6e88b 100644 --- a/posthog/api/team.py +++ b/posthog/api/team.py @@ -284,7 +284,11 @@ def create(self, validated_data: dict[str, Any], **kwargs) -> Team: # but ClickHouse doesn't support Saturday as the first day of the week, so we fall back to Sunday validated_data["week_start_day"] = 1 if week_start_day_for_user_ip_location == 1 else 0 - team = Team.objects.create_with_data(organization=self.context["view"].organization, **validated_data) + team = Team.objects.create_with_data( + initiating_user=self.context["request"].user, + organization=self.context["view"].organization, + **validated_data, + ) request.user.current_team = team request.user.team = request.user.current_team # Update cached property diff --git a/posthog/api/test/test_routing.py b/posthog/api/test/test_routing.py index 52712ae0500de0..24063d1a655e52 100644 --- a/posthog/api/test/test_routing.py +++ b/posthog/api/test/test_routing.py @@ -45,7 +45,9 @@ def setUp(self): super().setUp() other_org, _, other_org_team = Organization.objects.bootstrap(user=self.user) self.other_org_annotation = Annotation.objects.create(team=other_org_team, organization=other_org) - _, other_project_team = Project.objects.create_with_team(organization=self.organization) + _, other_project_team = Project.objects.create_with_team( + initiating_user=self.user, organization=self.organization + ) self.other_project_annotation = Annotation.objects.create( team=other_project_team, organization=self.organization ) diff --git a/posthog/api/test/test_team.py b/posthog/api/test/test_team.py index 0c69c2b6620d0e..3e4c3e48d2dad5 100644 --- a/posthog/api/test/test_team.py +++ b/posthog/api/test/test_team.py @@ -268,7 +268,7 @@ def test_delete_team_activity_log(self): self.organization_membership.level = OrganizationMembership.Level.ADMIN self.organization_membership.save() - team: Team = Team.objects.create_with_data(organization=self.organization) + team: Team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) response = self.client.delete(f"/api/environments/{team.id}") assert response.status_code == 204 @@ -345,7 +345,7 @@ def test_delete_team_own_second( self.organization_membership.level = OrganizationMembership.Level.ADMIN self.organization_membership.save() - team: Team = Team.objects.create_with_data(organization=self.organization) + team: Team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual(Team.objects.filter(organization=self.organization).count(), 2) @@ -384,7 +384,7 @@ def test_delete_bulky_postgres_data(self): self.organization_membership.level = OrganizationMembership.Level.ADMIN self.organization_membership.save() - team: Team = Team.objects.create_with_data(organization=self.organization) + team: Team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual(Team.objects.filter(organization=self.organization).count(), 2) @@ -434,7 +434,7 @@ def test_delete_batch_exports(self): self.organization_membership.level = OrganizationMembership.Level.ADMIN self.organization_membership.save() - team: Team = Team.objects.create_with_data(organization=self.organization) + team: Team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) destination_data = { "type": "S3", @@ -543,6 +543,9 @@ def test_update_primary_dashboard(self): self.assertEqual(response_data["primary_dashboard"], d.id) def test_cant_set_primary_dashboard_to_another_teams_dashboard(self): + self.team.primary_dashboard_id = None # Remove the default primary dashboard from the picture + self.team.save() + team_2 = Team.objects.create(organization=self.organization, name="Default project") d = Dashboard.objects.create(name="Test", team=team_2) @@ -564,6 +567,14 @@ def test_is_generating_demo_data(self): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(response.json(), {"is_generating_demo_data": False}) + @patch("posthog.tasks.demo_create_data.create_data_for_demo_team.delay") + def test_org_member_can_create_demo_project(self, mock_create_data_for_demo_team: MagicMock): + self.organization_membership.level = OrganizationMembership.Level.MEMBER + self.organization_membership.save() + response = self.client.post("/api/environments/", {"name": "Hedgebox", "is_demo": True}) + mock_create_data_for_demo_team.assert_called_once() + self.assertEqual(response.status_code, status.HTTP_201_CREATED) + @freeze_time("2022-02-08") def test_team_float_config_can_be_serialized_to_activity_log(self): # regression test since this isn't true by default diff --git a/posthog/demo/legacy/__init__.py b/posthog/demo/legacy/__init__.py index a62ca3f350e234..194eb3fcc05428 100644 --- a/posthog/demo/legacy/__init__.py +++ b/posthog/demo/legacy/__init__.py @@ -19,21 +19,6 @@ def demo_route(request: HttpRequest): return render_template("demo.html", request=request, context={"api_token": project_api_token}) -def create_demo_team(organization: Organization, *args) -> Team: - team = Team.objects.create_with_data( - default_dashboards=False, - organization=organization, - name=TEAM_NAME, - ingested_event=True, - completed_snippet_onboarding=True, - session_recording_opt_in=True, - is_demo=True, - ) - create_demo_data(team) - EventDefinition.objects.get_or_create(team=team, name="$pageview") - return team - - def create_demo_data(team: Team, dashboards=True): WebDataGenerator(team, n_people=40).create(dashboards=dashboards) AppDataGenerator(team, n_people=100).create(dashboards=dashboards) diff --git a/posthog/models/organization.py b/posthog/models/organization.py index e64f45ff8abc43..cf1cd5c26986a2 100644 --- a/posthog/models/organization.py +++ b/posthog/models/organization.py @@ -72,7 +72,9 @@ def bootstrap( with transaction.atomic(using=self.db): organization = Organization.objects.create(**kwargs) - _, team = Project.objects.create_with_team(organization=organization, team_fields=team_fields) + _, team = Project.objects.create_with_team( + initiating_user=user, organization=organization, team_fields=team_fields + ) organization_membership: Optional[OrganizationMembership] = None if user is not None: organization_membership = OrganizationMembership.objects.create( diff --git a/posthog/models/project.py b/posthog/models/project.py index bc006bf6546c3c..edd1fdff4edcc6 100644 --- a/posthog/models/project.py +++ b/posthog/models/project.py @@ -7,11 +7,13 @@ from posthog.models.utils import sane_repr if TYPE_CHECKING: - from .team import Team + from posthog.models import Team, User class ProjectManager(models.Manager): - def create_with_team(self, *, team_fields: Optional[dict] = None, **kwargs) -> tuple["Project", "Team"]: + def create_with_team( + self, *, team_fields: Optional[dict] = None, initiating_user: Optional["User"], **kwargs + ) -> tuple["Project", "Team"]: from .team import Team if team_fields is None: @@ -22,8 +24,12 @@ def create_with_team(self, *, team_fields: Optional[dict] = None, **kwargs) -> t with transaction.atomic(using=self.db): common_id = Team.objects.increment_id_sequence() project = cast("Project", self.create(id=common_id, **kwargs)) - team = Team.objects.create( - id=common_id, organization_id=project.organization_id, project=project, **team_fields + team = Team.objects.create_with_data( + id=common_id, + organization_id=project.organization_id, + project=project, + initiating_user=initiating_user, + **team_fields, ) return project, team diff --git a/posthog/models/team/team.py b/posthog/models/team/team.py index 016d08e64830a4..3bce7301ea23dd 100644 --- a/posthog/models/team/team.py +++ b/posthog/models/team/team.py @@ -1,7 +1,8 @@ import re from decimal import Decimal from functools import lru_cache -from typing import TYPE_CHECKING, Any, Optional, cast +from typing import TYPE_CHECKING, Optional, cast +from uuid import UUID from zoneinfo import ZoneInfo from django.core.cache import cache import posthoganalytics @@ -26,7 +27,7 @@ from posthog.models.filters.mixins.utils import cached_property from posthog.models.filters.utils import GroupTypeIndex from posthog.models.instance_setting import get_instance_setting -from posthog.models.organization import Organization +from posthog.models.organization import OrganizationMembership from posthog.models.signals import mutable_receiver from posthog.models.utils import ( UUIDClassicModel, @@ -66,7 +67,7 @@ class TeamManager(models.Manager): def get_queryset(self): return super().get_queryset().defer(*DEPRECATED_ATTRS) - def set_test_account_filters(self, organization: Optional[Any]) -> list: + def set_test_account_filters(self, organization_id: Optional[UUID]) -> list: filters = [ { "key": "$host", @@ -75,10 +76,12 @@ def set_test_account_filters(self, organization: Optional[Any]) -> list: "type": "event", } ] - if organization: - example_emails = organization.members.only("email") + if organization_id: + example_emails_raw = OrganizationMembership.objects.filter(organization_id=organization_id).values_list( + "user__email", flat=True + ) generic_emails = GenericEmails() - example_emails = [email.email for email in example_emails if not generic_emails.is_generic(email.email)] + example_emails = [email for email in example_emails_raw if not generic_emails.is_generic(email)] if len(example_emails) > 0: example_email = re.search(r"@[\w.]+", example_emails[0]) if example_email: @@ -88,18 +91,25 @@ def set_test_account_filters(self, organization: Optional[Any]) -> list: ] return filters - def create_with_data( - self, user: Any = None, default_dashboards: bool = True, *, organization: Organization, **kwargs - ) -> "Team": - kwargs["test_account_filters"] = self.set_test_account_filters(organization) - team = cast("Team", self.create(organization=organization, **kwargs)) - - # Create default dashboards (skipped for demo projects) - if default_dashboards: - dashboard = Dashboard.objects.db_manager(self.db).create(name="My App Dashboard", pinned=True, team=team) - create_dashboard_from_template("DEFAULT_APP", dashboard) - team.primary_dashboard = dashboard - team.save() + def create_with_data(self, *, initiating_user: Optional["User"], **kwargs) -> "Team": + team = cast("Team", self.create(**kwargs)) + + if kwargs.get("is_demo"): + if initiating_user is None: + raise ValueError("initiating_user must be provided when creating a demo team") + team.kick_off_demo_data_generation(initiating_user) + return team # Return quickly, as the demo data and setup will be created asynchronously + + team.test_account_filters = self.set_test_account_filters( + kwargs.get("organization_id") or kwargs["organization"].id + ) + + # Create default dashboards + dashboard = Dashboard.objects.db_manager(self.db).create(name="My App Dashboard", pinned=True, team=team) + create_dashboard_from_template("DEFAULT_APP", dashboard) + team.primary_dashboard = dashboard + + team.save() return team def create(self, **kwargs): @@ -478,7 +488,15 @@ def reset_token_and_save(self, *, user: "User", is_impersonated_session: bool): ) def get_is_generating_demo_data(self) -> bool: - return cache.get(f"is_generating_demo_data_{self.pk}") == "True" + cache_key = f"is_generating_demo_data_{self.id}" + return cache.get(cache_key) == "True" + + def kick_off_demo_data_generation(self, initiating_user: "User") -> None: + from posthog.tasks.demo_create_data import create_data_for_demo_team + + cache_key = f"is_generating_demo_data_{self.id}" + cache.set(cache_key, "True") # Create an item in the cache that we can use to see if the demo data is ready + create_data_for_demo_team.delay(self.id, initiating_user.id, cache_key) def all_users_with_access(self) -> QuerySet["User"]: from ee.models.explicit_team_membership import ExplicitTeamMembership diff --git a/posthog/models/test/test_project.py b/posthog/models/test/test_project.py index 1e2e0cef2167ac..1fd7434f90da3d 100644 --- a/posthog/models/test/test_project.py +++ b/posthog/models/test/test_project.py @@ -7,6 +7,7 @@ class TestProject(BaseTest): def test_create_project_with_team_no_team_fields(self): project, team = Project.objects.create_with_team( + initiating_user=self.user, organization=self.organization, name="Test project", ) @@ -24,6 +25,7 @@ def test_create_project_with_team_no_team_fields(self): def test_create_project_with_team_with_team_fields(self): project, team = Project.objects.create_with_team( + initiating_user=self.user, organization=self.organization, name="Test project", team_fields={"name": "Test team", "access_control": True}, @@ -42,6 +44,7 @@ def test_create_project_with_team_uses_team_id_sequence(self): expected_common_id = Team.objects.increment_id_sequence() + 1 project, team = Project.objects.create_with_team( + initiating_user=self.user, organization=self.organization, name="Test project", team_fields={"name": "Test team", "access_control": True}, @@ -64,6 +67,7 @@ def test_create_project_with_team_does_not_create_if_team_fails(self, mock_creat with self.assertRaises(Exception): Project.objects.create_with_team( + initiating_user=self.user, organization=self.organization, name="Test project", team_fields={"name": "Test team", "access_control": True}, diff --git a/posthog/models/user.py b/posthog/models/user.py index 621c1d36429a72..748533be437cd5 100644 --- a/posthog/models/user.py +++ b/posthog/models/user.py @@ -80,7 +80,9 @@ def bootstrap( if create_team: team = create_team(organization, user) else: - team = Team.objects.create_with_data(user=user, organization=organization, **(team_fields or {})) + team = Team.objects.create_with_data( + initiating_user=user, organization=organization, **(team_fields or {}) + ) user.join(organization=organization, level=OrganizationMembership.Level.OWNER) return organization, team, user diff --git a/posthog/test/base.py b/posthog/test/base.py index 451264bfc205be..ebb07f9270b8a0 100644 --- a/posthog/test/base.py +++ b/posthog/test/base.py @@ -108,6 +108,7 @@ def _setup_test_data(klass): klass.organization = Organization.objects.create(name=klass.CONFIG_ORGANIZATION_NAME) klass.project, klass.team = Project.objects.create_with_team( + initiating_user=klass.user, organization=klass.organization, team_fields={ "api_token": klass.CONFIG_API_TOKEN, diff --git a/posthog/test/test_team.py b/posthog/test/test_team.py index 076fc21e5fe345..b47eb70e016e2b 100644 --- a/posthog/test/test_team.py +++ b/posthog/test/test_team.py @@ -76,7 +76,7 @@ def test_team_has_expected_defaults(self): self.assertEqual(team.autocapture_exceptions_errors_to_ignore, None) def test_create_team_with_test_account_filters(self): - team = Team.objects.create_with_data(organization=self.organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual( team.test_account_filters, [ @@ -99,7 +99,7 @@ def test_create_team_with_test_account_filters(self): user = User.objects.create(email="test@gmail.com") organization = Organization.objects.create() organization.members.set([user]) - team = Team.objects.create_with_data(organization=organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual( team.test_account_filters, [ @@ -113,7 +113,7 @@ def test_create_team_with_test_account_filters(self): ) def test_create_team_sets_primary_dashboard(self): - team = Team.objects.create_with_data(organization=self.organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertIsInstance(team.primary_dashboard, Dashboard) # Ensure insights are created and linked @@ -139,7 +139,7 @@ def test_preinstalled_are_autoenabled(self, mock_get): def test_team_on_cloud_uses_feature_flag_to_determine_person_on_events(self, mock_feature_enabled): with self.is_cloud(True): with override_instance_config("PERSON_ON_EVENTS_ENABLED", False): - team = Team.objects.create_with_data(organization=self.organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual( team.person_on_events_mode, PersonsOnEventsMode.PERSON_ID_OVERRIDE_PROPERTIES_ON_EVENTS ) @@ -162,7 +162,7 @@ def test_team_on_cloud_uses_feature_flag_to_determine_person_on_events(self, moc def test_team_on_self_hosted_uses_instance_setting_to_determine_person_on_events(self, mock_feature_enabled): with self.is_cloud(False): with override_instance_config("PERSON_ON_EVENTS_V2_ENABLED", True): - team = Team.objects.create_with_data(organization=self.organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual( team.person_on_events_mode, PersonsOnEventsMode.PERSON_ID_OVERRIDE_PROPERTIES_ON_EVENTS ) @@ -171,7 +171,7 @@ def test_team_on_self_hosted_uses_instance_setting_to_determine_person_on_events assert args_list[0][0] != "persons-on-events-v2-reads-enabled" with override_instance_config("PERSON_ON_EVENTS_V2_ENABLED", False): - team = Team.objects.create_with_data(organization=self.organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) self.assertEqual(team.person_on_events_mode, PersonsOnEventsMode.DISABLED) for args_list in mock_feature_enabled.call_args_list: # It is ok if we check other feature flags, just not `persons-on-events-v2-reads-enabled` @@ -179,7 +179,7 @@ def test_team_on_self_hosted_uses_instance_setting_to_determine_person_on_events def test_each_team_gets_project_with_default_name_and_same_id(self): # Can be removed once environments are fully rolled out - team = Team.objects.create_with_data(organization=self.organization) + team = Team.objects.create_with_data(initiating_user=self.user, organization=self.organization) project = Project.objects.filter(id=team.id).first() @@ -188,7 +188,7 @@ def test_each_team_gets_project_with_default_name_and_same_id(self): def test_each_team_gets_project_with_custom_name_and_same_id(self): # Can be removed once environments are fully rolled out - team = Team.objects.create_with_data(organization=self.organization, name="Hogflix") + team = Team.objects.create_with_data(organization=self.organization, initiating_user=self.user, name="Hogflix") project = Project.objects.filter(id=team.id).first() @@ -203,7 +203,7 @@ def test_team_not_created_if_project_creation_fails(self, mock_create): initial_project_count = Project.objects.count() with self.assertRaises(Exception): - Team.objects.create_with_data(organization=self.organization, name="Hogflix") + Team.objects.create_with_data(organization=self.organization, initiating_user=self.user, name="Hogflix") self.assertEqual(Team.objects.count(), initial_team_count) self.assertEqual(Project.objects.count(), initial_project_count)