From a6a60360e7db5568235888e868eca1d92a0a24b3 Mon Sep 17 00:00:00 2001
From: Phani Raj <phani@posthog.com>
Date: Tue, 23 Jul 2024 10:59:19 -0500
Subject: [PATCH] guard for nil iteration_count when updating survey iterations

---
 posthog/api/survey.py           |  9 ++++++++-
 posthog/api/test/test_survey.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/posthog/api/survey.py b/posthog/api/survey.py
index 8c18fe6170032..47dc027fffd2f 100644
--- a/posthog/api/survey.py
+++ b/posthog/api/survey.py
@@ -379,7 +379,14 @@ def update(self, instance: Survey, validated_data):
             instance.targeting_flag.save()
 
         iteration_count = validated_data.get("iteration_count")
-        if instance.current_iteration is not None and instance.current_iteration > iteration_count > 0:
+        if iteration_count is None:
+            iteration_count = 0
+
+        if (
+            instance.current_iteration is not None
+            and iteration_count is not None
+            and instance.current_iteration > iteration_count > 0
+        ):
             raise serializers.ValidationError(
                 f"Cannot change survey recurrence to {iteration_count}, should be at least {instance.current_iteration}"
             )
diff --git a/posthog/api/test/test_survey.py b/posthog/api/test/test_survey.py
index 6fd90dfbfb26c..f0f685784c58a 100644
--- a/posthog/api/test/test_survey.py
+++ b/posthog/api/test/test_survey.py
@@ -2346,6 +2346,35 @@ def test_cannot_reduce_iterations_lt_current_iteration(self):
         assert response.status_code == status.HTTP_400_BAD_REQUEST
         assert response.json()["detail"] == "Cannot change survey recurrence to 1, should be at least 2"
 
+    def test_can_set_recurrence_if_previously_disabled(self):
+        survey = self._create_recurring_survey()
+        survey.current_iteration = 2
+        survey.save()
+        response = self.client.patch(
+            f"/api/projects/{self.team.id}/surveys/{survey.id}/",
+            data={
+                "start_date": datetime.now() - timedelta(days=1),
+            },
+        )
+        response_data = response.json()
+        survey.refresh_from_db()
+        survey.current_iteration = 2
+        survey.save()
+        response = self.client.patch(
+            f"/api/projects/{self.team.id}/surveys/{survey.id}/",
+            data={
+                "start_date": datetime.now() - timedelta(days=1),
+                "iteration_count": 2,
+                "iteration_frequency_days": 30,
+            },
+        )
+
+        assert response.status_code == status.HTTP_200_OK
+        response_data = response.json()
+        assert response_data["iteration_start_dates"] is not None
+        assert len(response_data["iteration_start_dates"]) == 2
+        assert response_data["current_iteration"] == 1
+
     def test_can_turn_off_recurring_schedule(self):
         survey = self._create_recurring_survey()
         response = self.client.patch(