Add a test to make sure funnel cohort persons is returned (#5103)

* add a test to make sure funnel cohort persons is returned

* delete comment

* checking persons on cohort return in original test

ee/clickhouse/queries/funnels/base.py

@@ -250,9 +250,12 @@ class ClickhouseFunnelBase(ABC, Funnel):
             conditions.append("steps = %(step_num)s")
 
         if self._filter.funnel_step_breakdown:
-            self.params.update(
-                {"breakdown_prop_value": [val.strip() for val in self._filter.funnel_step_breakdown.split(",")]}
+            prop_vals = (
+                [val.strip() for val in self._filter.funnel_step_breakdown.split(",")]
+                if isinstance(self._filter.funnel_step_breakdown, str)
+                else [self._filter.funnel_step_breakdown]
             )
+            self.params.update({"breakdown_prop_value": prop_vals})
             conditions.append("prop IN %(breakdown_prop_value)s")
 
         return " AND ".join(conditions)

ee/clickhouse/queries/funnels/test/breakdown_cases.py

@@ -697,7 +697,7 @@ def funnel_breakdown_test_factory(Funnel, FunnelPerson, _create_event, _create_p
 
         def test_funnel_cohort_breakdown(self):
             # This caused some issues with SQL parsing
-            _create_person(distinct_ids=[f"person1"], team_id=self.team.pk, properties={"key": "value"})
+            person = _create_person(distinct_ids=[f"person1"], team_id=self.team.pk, properties={"key": "value"})
             _create_event(
                 team=self.team,
                 event="sign up",
@@ -723,5 +723,7 @@ def funnel_breakdown_test_factory(Funnel, FunnelPerson, _create_event, _create_p
             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

ee/clickhouse/queries/funnels/test/test_funnel_persons.py

@@ -4,9 +4,8 @@ from ee.clickhouse.models.event import create_event
 from ee.clickhouse.queries.funnels.funnel import ClickhouseFunnel
 from ee.clickhouse.queries.funnels.funnel_persons import ClickhouseFunnelPersons
 from ee.clickhouse.util import ClickhouseTestMixin
-from posthog.constants import INSIGHT_FUNNELS, TRENDS_FUNNEL
-from posthog.models import Filter
-from posthog.models.filters.mixins.funnel import FunnelWindowDaysMixin
+from posthog.constants import INSIGHT_FUNNELS
+from posthog.models import Cohort, Filter
 from posthog.models.person import Person
 from posthog.test.base import APIBaseTest
 
@@ -226,3 +225,23 @@ class TestFunnelPersons(ClickhouseTestMixin, APIBaseTest):
             filter.with_data({"funnel_step_breakdown": "Safari, Chrome"}), self.team
         )._exec_query()
         self.assertCountEqual([val[0] for val in results], [person2.uuid, person1.uuid])
+
+    def test_funnel_cohort_breakdown_persons(self):
+        person = _create_person(distinct_ids=[f"person1"], team_id=self.team.pk, properties={"key": "value"})
+        _create_event(
+            team=self.team, event="sign up", distinct_id=f"person1", properties={}, timestamp="2020-01-02T12:00:00Z",
+        )
+        cohort = Cohort.objects.create(team=self.team, name="test_cohort", groups=[{"properties": {"key": "value"}}])
+        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,
+            "funnel_step": 1,
+            "breakdown_type": "cohort",
+            "breakdown": [cohort.pk],
+        }
+        filter = Filter(data=filters)
+        results = ClickhouseFunnelPersons(filter, self.team)._exec_query()
+        self.assertEqual(results[0][0], person.uuid)

posthog/models/filters/mixins/funnel.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import Optional
+from typing import Optional, Union
 
 from posthog.constants import (
     BIN_COUNT,
@@ -87,7 +87,7 @@ class FunnelPersonsStepMixin(BaseParamMixin):
 
 class FunnelPersonsStepBreakdownMixin(BaseParamMixin):
     @cached_property
-    def funnel_step_breakdown(self) -> Optional[str]:
+    def funnel_step_breakdown(self) -> Optional[Union[str, int]]:
         return self._data.get(FUNNEL_STEP_BREAKDOWN)
 
     @include_dict