From ac794aaca253b56b22573cc1e547595d57e97f3c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Jakes=CC=8C?= <jan@jakes.pro>
Date: Mon, 18 Nov 2019 16:43:53 +0100
Subject: [PATCH] Use processTaskStrategy() instead of custom logic in
 Migration worker

[MAILPOET-2538]
---
 lib/Cron/Workers/SendingQueue/Migration.php         | 10 +++-------
 lib/Cron/Workers/SimpleWorker.php                   | 13 +++++++++----
 .../Cron/Workers/SendingQueue/MigrationTest.php     |  9 +++++----
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/lib/Cron/Workers/SendingQueue/Migration.php b/lib/Cron/Workers/SendingQueue/Migration.php
index 1304023782..0d4f734d73 100644
--- a/lib/Cron/Workers/SendingQueue/Migration.php
+++ b/lib/Cron/Workers/SendingQueue/Migration.php
@@ -21,7 +21,7 @@ class Migration extends SimpleWorker {
     return empty($completed_tasks);
   }
 
-  function prepareTask(ScheduledTask $task) {
+  function prepareTaskStrategy(ScheduledTask $task) {
     $unmigrated_columns = $this->checkUnmigratedColumnsExist();
     $unmigrated_queues_count = 0;
     $unmigrated_queue_subscribers = [];
@@ -43,8 +43,7 @@ class Migration extends SimpleWorker {
 
     // pause sending while the migration is in process
     $this->pauseSending();
-
-    return parent::prepareTask($task);
+    return true;
   }
 
   function pauseSending() {
@@ -74,13 +73,10 @@ class Migration extends SimpleWorker {
     }
   }
 
-  function processTask(ScheduledTask $task) {
+  function processTaskStrategy(ScheduledTask $task) {
     $this->migrateSendingQueues();
     $this->migrateSubscribers();
-
-    $this->complete($task);
     $this->resumeSending();
-
     return true;
   }
 
diff --git a/lib/Cron/Workers/SimpleWorker.php b/lib/Cron/Workers/SimpleWorker.php
index 1efc896ad7..f24b4db91e 100644
--- a/lib/Cron/Workers/SimpleWorker.php
+++ b/lib/Cron/Workers/SimpleWorker.php
@@ -93,13 +93,14 @@ abstract class SimpleWorker {
   }
 
   function prepareTask(ScheduledTask $task) {
-    $task->status = null;
-    $task->save();
-
     // abort if execution limit is reached
     $this->cron_helper->enforceExecutionLimit($this->timer);
 
-    return true;
+    $prepare_completed = $this->prepareTaskStrategy($task);
+    if ($prepare_completed) {
+      $task->status = null;
+      $task->save();
+    }
   }
 
   function processTask(ScheduledTask $task) {
@@ -133,6 +134,10 @@ abstract class SimpleWorker {
     return (bool)$completed;
   }
 
+  function prepareTaskStrategy(ScheduledTask $task) {
+    return true;
+  }
+
   function processTaskStrategy(ScheduledTask $task) {
     return true;
   }
diff --git a/tests/integration/Cron/Workers/SendingQueue/MigrationTest.php b/tests/integration/Cron/Workers/SendingQueue/MigrationTest.php
index 71a267b320..9c50693518 100644
--- a/tests/integration/Cron/Workers/SendingQueue/MigrationTest.php
+++ b/tests/integration/Cron/Workers/SendingQueue/MigrationTest.php
@@ -15,6 +15,9 @@ use MailPoet\Tasks\Sending as SendingTask;
 use MailPoet\WP\Functions as WPFunctions;
 
 class MigrationTest extends \MailPoetTest {
+  /** @var Migration */
+  private $worker;
+
   function _before() {
     parent::_before();
     // Alter table to test migration
@@ -63,17 +66,15 @@ class MigrationTest extends \MailPoetTest {
     $this->worker->pauseSending();
     expect(MailerLog::isSendingPaused())->true();
     $task = $this->createScheduledTask();
-    $result = $this->worker->prepareTask($task);
-    expect($result)->false();
+    $this->worker->prepareTask($task);
     expect(MailerLog::isSendingPaused())->false();
   }
 
   function testItCompletesTaskIfThereIsNothingToMigrate() {
     SendingQueue::deleteMany();
     $task = $this->createScheduledTask();
-    $result = $this->worker->prepareTask($task);
+    $this->worker->prepareTask($task);
     expect(ScheduledTask::findOne($task->id)->status)->equals(ScheduledTask::STATUS_COMPLETED);
-    expect($result)->false();
   }
 
   function testItMigratesSendingQueuesAndSubscribers() {