some test fixes

Author: nickbest-ph

posthog/models/cohort/cohort.py

@@ -27,7 +27,6 @@
 from posthog.models.property import Property, PropertyGroup
 from posthog.models.utils import RootTeamManager, RootTeamMixin, sane_repr
 from posthog.person_db_router import PERSONS_DB_FOR_WRITE
-from posthog.personhog_client.gate import use_personhog
 from posthog.settings.base_variables import TEST
 
 if TYPE_CHECKING:
@@ -391,6 +390,8 @@ def _get_uuids_for_distinct_ids_batch(self, distinct_ids: list[str], team_id: in
         # You're going to be tempted to exclude people already in the cohort, but that's not only NOT
         # necessary, but it leads to query timeouts. The insert_users_list_by_uuid handles ensuring we
         # don't insert people that are already in the cohort efficiently.
+        from posthog.personhog_client.gate import use_personhog
+
         if use_personhog():
             persons = get_persons_by_distinct_ids(team_id, list(distinct_ids))
             return [str(person.uuid) for person in persons]

posthog/models/person/test/test_util_personhog_routing.py

@@ -997,15 +997,12 @@ def test_single_person_multiple_distinct_ids(self):
         result = get_persons_mapped_by_distinct_id(self.team.pk, ["did-1", "did-2"])
 
         # Both distinct_ids should map to the same person
+        assert len(result) == 2
         assert str(result["did-1"].uuid) == str(person.uuid)
+        assert str(result["did-2"].uuid) == str(person.uuid)
         # Each entry should carry only the distinct_id that was used as the key
         assert result["did-1"].distinct_ids == ["did-1"]
-        if self.personhog:
-            # Personhog path deduplicates by person_id, so only one entry
-            assert len(result) == 1
-        else:
-            assert len(result) == 2
-            assert result["did-2"].distinct_ids == ["did-2"]
+        assert result["did-2"].distinct_ids == ["did-2"]
 
     def test_multiple_persons(self):
         p1 = self._seed_person(team=self.team, distinct_ids=["alice"], properties={"name": "Alice"})

posthog/models/person/util.py

@@ -245,6 +245,16 @@ def _fetch_persons_by_distinct_ids_via_personhog(
             "personhog_team_mismatch", operation="get_persons_by_distinct_ids", team_id=team_id, dropped=mismatched
         )
 
+    # The RPC returns one result per distinct_id, so the same person can
+    # appear multiple times.  Deduplicate by person_id to return unique persons.
+    seen_person_ids: set[int] = set()
+    unique_results = []
+    for r in valid_results:
+        if r.person.id not in seen_person_ids:
+            seen_person_ids.add(r.person.id)
+            unique_results.append(r)
+    valid_results = unique_results
+
     person_ids = [r.person.id for r in valid_results]
     if not person_ids:
         return []

posthog/personhog_client/fake_client.py

@@ -208,14 +208,17 @@ def get_persons_by_distinct_ids_in_team(
     ) -> person_pb2.PersonsByDistinctIdsInTeamResponse:
         self.calls.append(_Call("get_persons_by_distinct_ids_in_team", request))
         results = []
-        seen_person_ids: set[int] = set()
         for did in request.distinct_ids:
             person = self._persons_by_distinct_id.get((request.team_id, did))
-            if person and person.id not in seen_person_ids:
-                seen_person_ids.add(person.id)
+            if person:
                 results.append(person_pb2.PersonWithDistinctIds(distinct_id=did, person=person))
         return person_pb2.PersonsByDistinctIdsInTeamResponse(results=results)
 
+    # NOTE: the real RPC returns one result per requested distinct_id (no
+    # deduplication by person).  Callers that need unique persons (e.g.
+    # _fetch_persons_by_distinct_ids_via_personhog) must deduplicate
+    # themselves.
+
     def get_distinct_ids_for_person(
         self, request: person_pb2.GetDistinctIdsForPersonRequest
     ) -> person_pb2.GetDistinctIdsForPersonResponse:

posthog/queries/actor_base_query.py

@@ -21,7 +21,6 @@
 from posthog.models.person.person import READ_DB_FOR_PERSONS
 from posthog.personhog_client.client import get_personhog_client
 from posthog.personhog_client.converters import proto_person_to_model
-from posthog.personhog_client.gate import use_personhog
 from posthog.personhog_client.metrics import (
     PERSONHOG_ROUTING_ERRORS_TOTAL,
     PERSONHOG_ROUTING_TOTAL,
@@ -308,6 +307,8 @@ def get_people(
     distinct_id_limit: int | None = 1000,
 ) -> tuple[Union[QuerySet[Person], list[Person]], list[SerializedPerson]]:
     """Get people from raw SQL results in data model and dict formats"""
+    from posthog.personhog_client.gate import use_personhog
+
     if use_personhog():
         try:
             persons = _fetch_people_via_personhog(team.pk, people_ids, distinct_id_limit)