diff --git a/core/modules/migrate/tests/src/Unit/MigrateTestCase.php b/core/modules/migrate/tests/src/Unit/MigrateTestCase.php
index b1f8dbc..a6a15b3 100644
--- a/core/modules/migrate/tests/src/Unit/MigrateTestCase.php
+++ b/core/modules/migrate/tests/src/Unit/MigrateTestCase.php
@@ -18,6 +18,11 @@
   protected $migrationConfiguration = array();
 
   /**
+   * @var \Drupal\migrate\Entity\MigrationInterface::STATUS_*
+   */
+  protected $migrationStatus;
+
+  /**
    * Retrieve a mocked migration.
    *
    * @return \Drupal\migrate\Entity\MigrationInterface
@@ -40,6 +45,20 @@ protected function getMigration() {
     $migration->expects($this->any())
       ->method('getIdMap')
       ->will($this->returnValue($this->idMap));
+
+    // We need the state to be toggled throughout the test so we store the value
+    // on the test class and use a return callback.
+    $migration->expects($this->any())
+      ->method('getMigrationStatus')
+      ->will($this->returnCallback(function() {
+        return $this->migrationStatus;
+      }));
+    $migration->expects($this->any())
+      ->method('setMigrationStatus')
+      ->will($this->returnCallback(function($status) {
+        $this->migrationStatus = $status;
+      }));
+
     $configuration = &$this->migrationConfiguration;
     $migration->expects($this->any())->method('get')->will($this->returnCallback(function ($argument) use (&$configuration) {
       return isset($configuration[$argument]) ? $configuration[$argument] : '';