fix types, add snapshot, clean test

posthog/tasks/stop_surveys_reached_target.py

@@ -1,14 +1,16 @@
 from itertools import groupby
-from django.db.models import UUIDField, DateTimeField, Q
+from django.db.models import Q
 from django.utils import timezone
+from uuid import UUID
+from datetime import datetime
 
 from posthog.clickhouse.client.connection import Workload
 from posthog.client import sync_execute
 from posthog.models import Survey
 
 
 def _get_surveys_response_counts(
-    surveys_ids: list[UUIDField], team_id: UUIDField, earliest_survey_creation_date: DateTimeField
+    surveys_ids: list[UUID], team_id: int, earliest_survey_creation_date: datetime
 ) -> dict[str, int]:
     data = sync_execute(
         """

posthog/tasks/test/__snapshots__/test_stop_surveys_reached_target.ambr

@@ -0,0 +1,27 @@
+# serializer version: 1
+# name: TestStopSurveysReachedTarget.test_stop_surveys_with_enough_responses
+  '''
+
+  SELECT JSONExtractString(properties, '$survey_id') as survey_id,
+         count()
+  FROM events
+  WHERE event = 'survey sent'
+    AND team_id = 2
+    AND timestamp >= '2021-12-29 20:00:00'
+    AND survey_id in ['017e12ef-9c00-0000-8f06-7166eda03e1f', '017e12ef-9c00-0001-ddc8-81795e948930']
+  GROUP BY survey_id
+  '''
+# ---
+# name: TestStopSurveysReachedTarget.test_stop_surveys_with_enough_responses.1
+  '''
+
+  SELECT JSONExtractString(properties, '$survey_id') as survey_id,
+         count()
+  FROM events
+  WHERE event = 'survey sent'
+    AND team_id = 2
+    AND timestamp >= '2022-01-01 00:00:00'
+    AND survey_id in ['017e12ef-9c00-0002-b9f3-4a4e9794e506']
+  GROUP BY survey_id
+  '''
+# ---

posthog/tasks/test/test_stop_surveys_reached_target.py

@@ -1,10 +1,11 @@
+from freezegun import freeze_time
 from posthog.models import Survey, Organization, Team, User, FeatureFlag
 from django.test import TestCase
 from typing import Optional
 from dateutil.relativedelta import relativedelta
 from datetime import timedelta, datetime
 from django.utils.timezone import now
-from posthog.test.base import _create_event, flush_persons_and_events, ClickhouseTestMixin
+from posthog.test.base import _create_event, flush_persons_and_events, ClickhouseTestMixin, snapshot_clickhouse_queries
 from posthog.tasks.stop_surveys_reached_target import stop_surveys_reached_target
 
 
@@ -38,6 +39,8 @@ def _create_event_for_survey(
             team=survey.team,
         )
 
+    @freeze_time("2022-01-01")
+    @snapshot_clickhouse_queries
     def test_stop_surveys_with_enough_responses(self) -> None:
         surveys = [
             Survey.objects.create(
@@ -46,26 +49,27 @@ def test_stop_surveys_with_enough_responses(self) -> None:
                 created_by=self.user,
                 linked_flag=self.flag,
                 responses_limit=1,
-                created_at=now(),
             ),
             Survey.objects.create(
                 name="2",
                 team=self.team1,
                 created_by=self.user,
                 linked_flag=self.flag,
                 responses_limit=1,
-                created_at=now(),
             ),
             Survey.objects.create(
                 name="3",
                 team=self.team2,
                 created_by=self.user,
                 linked_flag=self.flag,
                 responses_limit=1,
-                created_at=now(),
             ),
         ]
 
+        # TRICKY: We can't override created at in the model create because of auto_now_add, so we need to update it manually
+        surveys[1].created_at = now() - relativedelta(days=2, hours=4)
+        surveys[1].save()
+
         for survey in surveys:
             self._create_event_for_survey(survey)
 
@@ -158,8 +162,10 @@ def test_do_not_stop_already_stopped_survey_with_responses_limit(self) -> None:
             linked_flag=self.flag,
             responses_limit=1,
             end_date=now() - relativedelta(hours=1),
-            created_at=now() - relativedelta(hours=1),
         )
+        survey.created_at = now() - relativedelta(hours=1)
+        survey.save()
+
         self._create_event_for_survey(survey, event="survey sent")
         flush_persons_and_events()