From bcffcde7839c13fec1351a31720aabecc593f516 Mon Sep 17 00:00:00 2001
From: Eric Duong <eeoneric@gmail.com>
Date: Tue, 13 Jul 2021 11:15:12 -0400
Subject: [PATCH] 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         |  7 ++++--
 .../queries/funnels/test/breakdown_cases.py   |  4 ++-
 .../funnels/test/test_funnel_persons.py       | 25 ++++++++++++++++---
 posthog/models/filters/mixins/funnel.py       |  4 +--
 4 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/ee/clickhouse/queries/funnels/base.py b/ee/clickhouse/queries/funnels/base.py
index 78ce97802cc..3ba4341bede 100644
--- a/ee/clickhouse/queries/funnels/base.py
+++ b/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)
diff --git a/ee/clickhouse/queries/funnels/test/breakdown_cases.py b/ee/clickhouse/queries/funnels/test/breakdown_cases.py
index fc5f44754c2..764c4e002b0 100644
--- a/ee/clickhouse/queries/funnels/test/breakdown_cases.py
+++ b/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
diff --git a/ee/clickhouse/queries/funnels/test/test_funnel_persons.py b/ee/clickhouse/queries/funnels/test/test_funnel_persons.py
index acf25dc208a..e0308f3fbd0 100644
--- a/ee/clickhouse/queries/funnels/test/test_funnel_persons.py
+++ b/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)
diff --git a/posthog/models/filters/mixins/funnel.py b/posthog/models/filters/mixins/funnel.py
index 80ad2a64363..09197989af4 100644
--- a/posthog/models/filters/mixins/funnel.py
+++ b/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