def test_insert_by_distinct_id_or_email(self): Person.objects.create(team_id=self.team.pk, properties={"email": "*****@*****.**"}, distinct_ids=["1"]) Person.objects.create(team_id=self.team.pk, distinct_ids=["123"]) Person.objects.create(team_id=self.team.pk, distinct_ids=["2"]) # Team leakage team2 = Team.objects.create(organization=self.organization) Person.objects.create(team=team2, properties={"email": "*****@*****.**"}) cohort = Cohort.objects.create(team=self.team, groups=[], is_static=True) cohort.insert_users_by_list(["*****@*****.**", "123"]) cohort = Cohort.objects.get() results = get_person_ids_by_cohort_id(self.team, cohort.id) self.assertEqual(len(results), 2) self.assertEqual(cohort.is_calculating, False) # test SQLi Person.objects.create(team_id=self.team.pk, distinct_ids=["'); truncate person_static_cohort; --"]) cohort.insert_users_by_list(["'); truncate person_static_cohort; --", "123"]) results = sync_execute("select count(1) from person_static_cohort")[0][0] self.assertEqual(results, 3) # If we accidentally call calculate_people it shouldn't erase people cohort.calculate_people() results = get_person_ids_by_cohort_id(self.team, cohort.id) self.assertEqual(len(results), 3) # if we add people again, don't increase the number of people in cohort cohort.insert_users_by_list(["123"]) results = get_person_ids_by_cohort_id(self.team, cohort.id) self.assertEqual(len(results), 3)
def _clickhouse_persons_query(self, batch_size=10000, offset=0): from ee.clickhouse.models.cohort import get_person_ids_by_cohort_id uuids = get_person_ids_by_cohort_id(team=self.team, cohort_id=self.pk, limit=batch_size, offset=offset) return Person.objects.filter(uuid__in=uuids, team=self.team)
def test_cohort_get_person_ids_by_cohort_id(self): user1 = _create_person(distinct_ids=["user1"], team_id=self.team.pk, properties={"$some_prop": "something"}) user2 = _create_person(distinct_ids=["user2"], team_id=self.team.pk, properties={"$some_prop": "another"}) user3 = _create_person(distinct_ids=["user3"], team_id=self.team.pk, properties={"$some_prop": "something"}) cohort = Cohort.objects.create( team=self.team, groups=[{"properties": {"$some_prop": "something"}}], name="cohort1", ) results = get_person_ids_by_cohort_id(self.team, cohort.id) self.assertEqual(len(results), 2) self.assertIn(user1.uuid, results) self.assertIn(user3.uuid, results)
def _clickhouse_persons_query(self): from ee.clickhouse.models.cohort import get_person_ids_by_cohort_id uuids = get_person_ids_by_cohort_id(team=self.team, cohort_id=self.pk) return Person.objects.filter(uuid__in=uuids, team=self.team)