Fix single cohort bug (#5225)

* fix single cohort bug * fix test name
PostHog · Jul 20, 2021 · 993278a · 993278a
1 parent 8b6b6ae
commit 993278a
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 2 deletions.
diff --git a/ee/clickhouse/queries/breakdown_props.py b/ee/clickhouse/queries/breakdown_props.py
@@ -127,10 +127,14 @@ def _format_all_query(team_id: int, filter: Filter, **kwargs) -> Tuple[str, Dict
 
 def format_breakdown_cohort_join_query(team_id: int, filter: Filter, **kwargs) -> Tuple[str, List, Dict]:
     entity = kwargs.pop("entity", None)
-    cohorts = Cohort.objects.filter(team_id=team_id, pk__in=[b for b in filter.breakdown if b != "all"])
+    cohorts = (
+        Cohort.objects.filter(team_id=team_id, pk__in=[b for b in filter.breakdown if b != "all"])
+        if isinstance(filter.breakdown, list)
+        else Cohort.objects.filter(team_id=team_id, pk=filter.breakdown)
+    )
     cohort_queries, params = _parse_breakdown_cohorts(cohorts)
     ids = [cohort.pk for cohort in cohorts]
-    if "all" in filter.breakdown:
+    if isinstance(filter.breakdown, list) and "all" in filter.breakdown:
         all_query, all_params = _format_all_query(team_id, filter, entity=entity)
         cohort_queries.append(all_query)
         params = {**params, **all_params}

diff --git a/ee/clickhouse/queries/funnels/test/breakdown_cases.py b/ee/clickhouse/queries/funnels/test/breakdown_cases.py
@@ -726,4 +726,23 @@ def test_funnel_cohort_breakdown(self):
             self.assertCountEqual(self._get_people_at_step(filter, 1, cohort.pk), [person.uuid])
             self.assertCountEqual(self._get_people_at_step(filter, 2, cohort.pk), [])
 
+            # non array
+            filters = {
+                "events": [{"id": "sign up", "order": 0}, {"id": "play movie", "order": 1}, {"id": "buy", "order": 2},],
+                "insight": INSIGHT_FUNNELS,
+                "date_from": "2020-01-01",
+                "date_to": "2020-01-08",
+                "funnel_window_days": 7,
+                "breakdown_type": "cohort",
+                "breakdown": cohort.pk,
+            }
+            filter = Filter(data=filters)
+            funnel = ClickhouseFunnel(filter, self.team)
+
+            result = funnel.run()
+            self.assertEqual(len(result[0]), 3)
+            self.assertEqual(result[0][0]["breakdown"], "test_cohort")
+            self.assertCountEqual(self._get_people_at_step(filter, 1, cohort.pk), [person.uuid])
+            self.assertCountEqual(self._get_people_at_step(filter, 2, cohort.pk), [])
+
     return TestFunnelBreakdown
diff --git a/ee/clickhouse/queries/test/test_trends.py b/ee/clickhouse/queries/test/test_trends.py
@@ -666,3 +666,53 @@ def test_breakdown_multiple_cohorts(self):
 
         self.assertEqual(res[0]["count"], 1)
         self.assertEqual(res[1]["count"], 2)
+
+    def test_breakdown_single_cohort(self):
+        p1 = Person.objects.create(team_id=self.team.pk, distinct_ids=["p1"], properties={"key": "value"})
+        _create_event(
+            team=self.team,
+            event="$pageview",
+            distinct_id="p1",
+            timestamp="2020-01-02T12:00:00Z",
+            properties={"key": "val"},
+        )
+
+        p2 = Person.objects.create(team_id=self.team.pk, distinct_ids=["p2"], properties={"key_2": "value_2"})
+        _create_event(
+            team=self.team,
+            event="$pageview",
+            distinct_id="p2",
+            timestamp="2020-01-02T12:00:00Z",
+            properties={"key": "val"},
+        )
+
+        p3 = Person.objects.create(team_id=self.team.pk, distinct_ids=["p3"], properties={"key_2": "value_2"})
+        _create_event(
+            team=self.team,
+            event="$pageview",
+            distinct_id="p3",
+            timestamp="2020-01-02T12:00:00Z",
+            properties={"key": "val"},
+        )
+
+        cohort1 = _create_cohort(team=self.team, name="cohort_1", groups=[{"properties": {"key": "value"}}])
+
+        cohort1.calculate_people()
+        cohort1.calculate_people_ch()
+
+        with self.settings(USE_PRECALCULATED_CH_COHORT_PEOPLE=True):  # Normally this is False in tests
+            with freeze_time("2020-01-04T13:01:01Z"):
+                res = ClickhouseTrends().run(
+                    Filter(
+                        data={
+                            "date_from": "-7d",
+                            "events": [{"id": "$pageview"}],
+                            "properties": [],
+                            "breakdown": cohort1.pk,
+                            "breakdown_type": "cohort",
+                        }
+                    ),
+                    self.team,
+                )
+
+        self.assertEqual(res[0]["count"], 1)