Refactor regenerate_failed_zim tests to reduce Db cursor and duplication

Author: ziaddevv

wp1/logic/builder_test.py

@@ -273,6 +273,48 @@ def _insert_selection(
                 )
         self.wp10db.commit()
 
+    def _setup_failed_zim_regeneration_scenario(
+        self,
+        zim_schedule_id=b"schedule-123",
+        old_task_id="old-task-id",
+        create_new_selection=False,
+    ):
+        self._insert_zim_schedule(zim_schedule_id, self.builder.b_id)
+
+        with self.wp10db.cursor() as cursor:
+            # Insert selection v1 with failed ZIM
+            cursor.execute(
+                """INSERT INTO selections
+                (s_id, s_builder_id, s_updated_at, s_content_type, s_version,
+                 s_object_key, s_article_count)
+                VALUES (%s, %s, '20250102000000', 'text/tab-separated-values', 1,
+                        'old.tsv', 100)""",
+                (1, self.builder.b_id),
+            )
+            cursor.execute(
+                """INSERT INTO zim_tasks
+                (z_selection_id, z_zim_schedule_id, z_status, z_task_id)
+                VALUES (%s, %s, 'FAILED', %s)""",
+                (1, zim_schedule_id, old_task_id),
+            )
+            # Optionally create selection v2
+            if create_new_selection:
+                cursor.execute(
+                    """INSERT INTO selections
+                    (s_id, s_builder_id, s_updated_at, s_content_type, s_version,
+                     s_object_key, s_article_count)
+                    VALUES (%s, %s, '20250103000000', 'text/tab-separated-values', 2,
+                            'new.tsv', 100)""",
+                    (2, self.builder.b_id),
+                )
+                cursor.execute(
+                    "UPDATE builders SET b_current_version = 2 WHERE b_id = %s",
+                    (self.builder.b_id,),
+                )
+
+        self.wp10db.commit()
+        return zim_schedule_id
+
     def _get_builder_by_user_id(self):
         with self.wp10db.cursor() as cursor:
             cursor.execute(
@@ -1805,40 +1847,9 @@ def test_regenerate_zim_updates_old_task_when_selection_version_changed(
         Ensure the existing zim_task is updated (not dupplicated) when the selection version changes.
         """
         self._insert_builder()
-
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO selections
-                (s_id, s_builder_id, s_updated_at, s_content_type, s_version, s_object_key, s_article_count)
-                VALUES (%s, %s, '20250102000000', 'text/tab-separated-values', 1, 'old.tsv', 100)""",
-                (1, self.builder.b_id),
-            )
-        self.wp10db.commit()
-
-        zim_schedule_id = b"schedule-123"
-        self._insert_zim_schedule(zim_schedule_id, self.builder.b_id)
-
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO zim_tasks
-                (z_selection_id, z_zim_schedule_id, z_status, z_task_id)
-                VALUES (%s, %s, 'FAILED', 'old-task-id')""",
-                (1, zim_schedule_id),
-            )
-        self.wp10db.commit()
-
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO selections
-                (s_id, s_builder_id, s_updated_at, s_content_type, s_version, s_object_key, s_article_count)
-                VALUES (%s, %s, '20250103000000', 'text/tab-separated-values', 2, 'new.tsv', 100)""",
-                (2, self.builder.b_id),
-            )
-            cursor.execute(
-                "UPDATE builders SET b_current_version = 2 WHERE b_id = %s",
-                (self.builder.b_id,),
-            )
-        self.wp10db.commit()
+        zim_schedule_id = self._setup_failed_zim_regeneration_scenario(
+            create_new_selection=True
+        )
 
         mock_request_zimfarm_task.return_value = "new-task-id"
 
@@ -1850,11 +1861,11 @@ def test_regenerate_zim_updates_old_task_when_selection_version_changed(
         with self.wp10db.cursor() as cursor:
             cursor.execute("SELECT COUNT(*) as count FROM zim_tasks")
             count = cursor.fetchone()["count"]
-            self.assertEqual(1, count)
-
             cursor.execute("SELECT z_selection_id FROM zim_tasks")
             row = cursor.fetchone()
-            self.assertEqual(b"2", row["z_selection_id"])
+
+        self.assertEqual(1, count)
+        self.assertEqual(b"2", row["z_selection_id"])
 
     @patch("wp1.logic.builder.zimfarm.request_zimfarm_task")
     @patch(
@@ -1867,22 +1878,10 @@ def test_regenerate_zim_saves_new_task_id(
         test that new task_id from Zimfarm is saved correctly.
         """
         self._insert_builder()
-        self._insert_selection(
-            1, "text/tab-separated-values", builder_id=self.builder.b_id
+        zim_schedule_id = self._setup_failed_zim_regeneration_scenario(
+            old_task_id="task_v1"
         )
 
-        zim_schedule_id = b"schedule-123"
-        self._insert_zim_schedule(zim_schedule_id, self.builder.b_id)
-
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO zim_tasks
-                (z_selection_id, z_zim_schedule_id, z_status, z_task_id)
-                VALUES (%s, %s, 'FAILED', 'task_v1')""",
-                (1, zim_schedule_id),
-            )
-        self.wp10db.commit()
-
         mock_request_zimfarm_task.return_value = "task_v2"
 
         redis = MagicMock()
@@ -1897,7 +1896,8 @@ def test_regenerate_zim_saves_new_task_id(
         with self.wp10db.cursor() as cursor:
             cursor.execute("SELECT z_task_id FROM zim_tasks WHERE z_selection_id = 1")
             row = cursor.fetchone()
-            self.assertEqual(b"task_v2", row["z_task_id"])
+
+        self.assertEqual(b"task_v2", row["z_task_id"])
 
     @patch("wp1.logic.builder.zimfarm.request_zimfarm_task")
     @patch(
@@ -1909,51 +1909,19 @@ def test_regenerate_zim_updates_b_selection_zim_version(
         """
         test that b_selection_zim_version is updated for downloads to work.
         """
-        self._insert_builder(zim_version=1)  # startts with version 1
-
-        # selection v1
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO selections
-                (s_id, s_builder_id, s_updated_at, s_content_type, s_version, s_object_key, s_article_count)
-                VALUES (%s, %s, '20250102000000', 'text/tab-separated-values', 1, 'old.tsv', 100)""",
-                (1, self.builder.b_id),
-            )
-        self.wp10db.commit()
-
-        zim_schedule_id = b"schedule-123"
-        self._insert_zim_schedule(zim_schedule_id, self.builder.b_id)
-
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO zim_tasks
-                (z_selection_id, z_zim_schedule_id, z_status, z_task_id)
-                VALUES (%s, %s, 'FAILED', 'old-task')""",
-                (1, zim_schedule_id),
-            )
-        self.wp10db.commit()
-
-        # Selection v2
-        with self.wp10db.cursor() as cursor:
-            cursor.execute(
-                """INSERT INTO selections
-                (s_id, s_builder_id, s_updated_at, s_content_type, s_version, s_object_key, s_article_count)
-                VALUES (%s, %s, '20250103000000', 'text/tab-separated-values', 2, 'new.tsv', 100)""",
-                (2, self.builder.b_id),
-            )
-            cursor.execute(
-                "UPDATE builders SET b_current_version = 2 WHERE b_id = %s",
-                (self.builder.b_id,),
-            )
-        self.wp10db.commit()
+        self._insert_builder(zim_version=1)
+        zim_schedule_id = self._setup_failed_zim_regeneration_scenario(
+            old_task_id="old-task", create_new_selection=True
+        )
 
         with self.wp10db.cursor() as cursor:
             cursor.execute(
                 "SELECT b_selection_zim_version FROM builders WHERE b_id = %s",
                 (self.builder.b_id,),
             )
             version_before = cursor.fetchone()["b_selection_zim_version"]
-            self.assertEqual(1, version_before)
+
+        self.assertEqual(1, version_before)
 
         mock_request_zimfarm_task.return_value = "task_v2"
 
@@ -1968,7 +1936,8 @@ def test_regenerate_zim_updates_b_selection_zim_version(
                 (self.builder.b_id,),
             )
             version_after = cursor.fetchone()["b_selection_zim_version"]
-            self.assertEqual(2, version_after)
+
+        self.assertEqual(2, version_after)
 
     @patch("wp1.logic.builder.zimfarm.zim_file_url_for_task_id")
     def test_download_url_after_regeneration(self, mock_zim_file_url):
@@ -1978,6 +1947,7 @@ def test_download_url_after_regeneration(self, mock_zim_file_url):
         self._insert_builder()
 
         with self.wp10db.cursor() as cursor:
+            # Selection v1 with failed ZIM
             cursor.execute(
                 """INSERT INTO selections
                 (s_id, s_builder_id, s_updated_at, s_content_type, s_version, s_object_key, s_article_count)
@@ -1990,9 +1960,7 @@ def test_download_url_after_regeneration(self, mock_zim_file_url):
                 VALUES (%s, 'FAILED', 'old-failed-task')""",
                 (1,),
             )
-        self.wp10db.commit()
-
-        with self.wp10db.cursor() as cursor:
+            # Selection v2 with successful ZIM
             cursor.execute(
                 """INSERT INTO selections
                 (s_id, s_builder_id, s_updated_at, s_content_type, s_version, s_object_key, s_article_count)