diff --git a/posthog/tasks/test/test_update_cache.py b/posthog/tasks/test/test_update_cache.py index e23e5db105f40..1812d9d83a229 100644 --- a/posthog/tasks/test/test_update_cache.py +++ b/posthog/tasks/test/test_update_cache.py @@ -120,6 +120,119 @@ def test_update_cache_item_calls_right_class(self) -> None: self.assertEqual(updated_dashboard_item.refreshing, False) self.assertEqual(updated_dashboard_item.last_refresh, now()) + @freeze_time("2012-01-15") + @patch("posthog.tasks.update_cache.Funnel") + def test_update_cache_item_calls_right_funnel_class(self, funnel_mock: MagicMock) -> None: + #  basic funnel + filter = Filter( + data={ + "insight": "FUNNELS", + "events": [ + {"id": "$pageview", "order": 0, "type": "events"}, + {"id": "$pageview", "order": 1, "type": "events"}, + ], + } + ) + dashboard_item = self._create_dashboard(filter) + + funnel_mock.return_value.run.return_value = {} + with self.settings(EE_AVAILABLE=False): + update_cache_item( + generate_cache_key("{}_{}".format(filter.toJSON(), self.team.pk)), + CacheType.FUNNEL, + {"filter": filter.toJSON(), "team_id": self.team.pk,}, + ) + + updated_dashboard_item = DashboardItem.objects.get(pk=dashboard_item.pk) + self.assertEqual(updated_dashboard_item.refreshing, False) + self.assertEqual(updated_dashboard_item.last_refresh, now()) + funnel_mock.assert_called_once() + + @freeze_time("2012-01-15") + @patch("posthog.tasks.update_cache.ClickhouseFunnelUnordered", create=True) + @patch("posthog.tasks.update_cache.ClickhouseFunnelStrict", create=True) + @patch("posthog.tasks.update_cache.ClickhouseFunnelTimeToConvert", create=True) + @patch("posthog.tasks.update_cache.ClickhouseFunnelTrends", create=True) + @patch("posthog.tasks.update_cache.ClickhouseFunnel", create=True) + def test_update_cache_item_calls_right_funnel_class_clickhouse( + self, + funnel_mock: MagicMock, + funnel_trends_mock: MagicMock, + funnel_time_to_convert_mock: MagicMock, + funnel_strict_mock: MagicMock, + funnel_unordered_mock: MagicMock, + ) -> None: + #  basic funnel + base_filter = Filter( + data={ + "insight": "FUNNELS", + "events": [ + {"id": "$pageview", "order": 0, "type": "events"}, + {"id": "$pageview", "order": 1, "type": "events"}, + ], + } + ) + + with self.settings(EE_AVAILABLE=True, PRIMARY_DB="clickhouse"): + filter = base_filter + funnel_mock.return_value.run.return_value = {} + update_cache_item( + generate_cache_key("{}_{}".format(filter.toJSON(), self.team.pk)), + CacheType.FUNNEL, + {"filter": filter.toJSON(), "team_id": self.team.pk,}, + ) + funnel_mock.assert_called_once() + + # trends funnel + filter = base_filter.with_data({"funnel_viz_type": "trends"}) + funnel_trends_mock.return_value.run.return_value = {} + update_cache_item( + generate_cache_key("{}_{}".format(filter.toJSON(), self.team.pk)), + CacheType.FUNNEL, + {"filter": filter.toJSON(), "team_id": self.team.pk,}, + ) + + funnel_trends_mock.assert_called_once() + self.assertEqual(funnel_trends_mock.call_args[1]["funnel_order_class"], funnel_mock) + funnel_trends_mock.reset_mock() + + # trends unordered funnel + filter = base_filter.with_data({"funnel_viz_type": "trends", "funnel_order_type": "unordered"}) + funnel_trends_mock.return_value.run.return_value = {} + update_cache_item( + generate_cache_key("{}_{}".format(filter.toJSON(), self.team.pk)), + CacheType.FUNNEL, + {"filter": filter.toJSON(), "team_id": self.team.pk,}, + ) + + funnel_trends_mock.assert_called_once() + self.assertEqual(funnel_trends_mock.call_args[1]["funnel_order_class"], funnel_unordered_mock) + funnel_trends_mock.reset_mock() + + # time to convert strict funnel + filter = base_filter.with_data({"funnel_viz_type": "time_to_convert", "funnel_order_type": "strict"}) + funnel_time_to_convert_mock.return_value.run.return_value = {} + update_cache_item( + generate_cache_key("{}_{}".format(filter.toJSON(), self.team.pk)), + CacheType.FUNNEL, + {"filter": filter.toJSON(), "team_id": self.team.pk,}, + ) + + funnel_time_to_convert_mock.assert_called_once() + self.assertEqual(funnel_time_to_convert_mock.call_args[1]["funnel_order_class"], funnel_strict_mock) + funnel_time_to_convert_mock.reset_mock() + + # strict funnel + filter = base_filter.with_data({"funnel_order_type": "strict"}) + funnel_strict_mock.return_value.run.return_value = {} + update_cache_item( + generate_cache_key("{}_{}".format(filter.toJSON(), self.team.pk)), + CacheType.FUNNEL, + {"filter": filter.toJSON(), "team_id": self.team.pk,}, + ) + + funnel_strict_mock.assert_called_once() + def _test_refresh_dashboard_cache_types( self, filter: FilterType, cache_type: CacheType, patch_update_cache_item: MagicMock, ) -> None: diff --git a/posthog/tasks/update_cache.py b/posthog/tasks/update_cache.py index d2a0c75e777ba..ed849cb409c87 100644 --- a/posthog/tasks/update_cache.py +++ b/posthog/tasks/update_cache.py @@ -21,6 +21,7 @@ INSIGHT_TRENDS, TRENDS_LINEAR, TRENDS_STICKINESS, + FunnelOrderType, FunnelVizType, ) from posthog.decorators import CacheType @@ -40,9 +41,14 @@ from ee.clickhouse.queries.clickhouse_paths import ClickhousePaths from ee.clickhouse.queries.clickhouse_retention import ClickhouseRetention from ee.clickhouse.queries.clickhouse_stickiness import ClickhouseStickiness - from ee.clickhouse.queries.funnels.funnel import ClickhouseFunnel - from ee.clickhouse.queries.funnels.funnel_time_to_convert import ClickhouseFunnelTimeToConvert - from ee.clickhouse.queries.funnels.funnel_trends import ClickhouseFunnelTrends + from ee.clickhouse.queries.funnels import ( + ClickhouseFunnel, + ClickhouseFunnelBase, + ClickhouseFunnelStrict, + ClickhouseFunnelTimeToConvert, + ClickhouseFunnelTrends, + ClickhouseFunnelUnordered, + ) from ee.clickhouse.queries.sessions.clickhouse_sessions import ClickhouseSessions from ee.clickhouse.queries.trends.clickhouse_trends import ClickhouseTrends @@ -165,17 +171,25 @@ def _calculate_funnel(filter: Filter, key: str, team_id: int) -> List[Dict[str, dashboard_items = DashboardItem.objects.filter(team_id=team_id, filters_hash=key) dashboard_items.update(refreshing=True) - insight_class: Union[Type[Funnel]] + team = Team(pk=team_id) + if is_clickhouse_enabled(): + funnel_order_class: Type[ClickhouseFunnelBase] = ClickhouseFunnel + if filter.funnel_order_type == FunnelOrderType.UNORDERED: + funnel_order_class = ClickhouseFunnelUnordered + elif filter.funnel_order_type == FunnelOrderType.STRICT: + funnel_order_class = ClickhouseFunnelStrict + if filter.funnel_viz_type == FunnelVizType.TRENDS: - insight_class = ClickhouseFunnelTrends + result = ClickhouseFunnelTrends(team=team, filter=filter, funnel_order_class=funnel_order_class).run() elif filter.funnel_viz_type == FunnelVizType.TIME_TO_CONVERT: - insight_class = ClickhouseFunnelTimeToConvert + result = ClickhouseFunnelTimeToConvert( + team=team, filter=filter, funnel_order_class=funnel_order_class + ).run() else: - insight_class = ClickhouseFunnel + result = funnel_order_class(team=team, filter=filter).run() else: - insight_class = Funnel + result = Funnel(filter=filter, team=team).run() - result = insight_class(filter=filter, team=Team(pk=team_id)).run() dashboard_items.update(last_refresh=timezone.now(), refreshing=False) return result