diff --git a/src/Tests/MigrateUpgradeTestBase.php b/src/Tests/MigrateUpgradeTestBase.php
index b5f0da5..527b3d8 100644
--- a/src/Tests/MigrateUpgradeTestBase.php
+++ b/src/Tests/MigrateUpgradeTestBase.php
@@ -106,6 +106,11 @@ abstract class MigrateUpgradeTestBase extends WebTestBase {
    */
   protected function testMigrateUpgrade() {
     $connection_options = $this->sourceDatabase->getConnectionOptions();
+
+    foreach (array_keys(\Drupal::entityTypeManager()->getDefinitions()) as $entity_type) {
+      $original_counts[$entity_type] = count(\Drupal::entityTypeManager()->getStorage($entity_type)->loadMultiple());
+    }
+
     $this->drupalGet('/upgrade');
     $this->assertText('Upgrade a Drupal site by importing it into a clean and empty new install of Drupal 8. You will lose any existing configuration once you import your site into it. See the upgrading handbook for more detailed information.');
 
@@ -143,6 +148,23 @@ abstract class MigrateUpgradeTestBase extends WebTestBase {
       $expected_count = isset($expected_counts[$entity_type]) ? $expected_counts[$entity_type] : 0;
       $this->assertEqual($expected_count, $real_count, "Found $real_count $entity_type entities, expected $expected_count.");
     }
+
+    // Test the rollback process.
+    $this->drupalGet('/upgrade');
+    $this->assertText('An upgrade has already been performed on this site.');
+
+    $edits = ['upgrade_option' => 2];
+    $this->drupalPostForm(NULL, $edits, t('Continue'));
+    $this->assertText('All previously-imported content, as well as configuration such as field definitions, will be removed.');
+
+    $this->drupalPostForm(NULL, [], t('Perform rollback'));
+    $this->assertText('Rollback of the upgrade is complete - you may now start the upgrade process from scratch.');
+
+    foreach (array_keys(\Drupal::entityTypeManager()->getDefinitions()) as $entity_type) {
+      $real_count = count(\Drupal::entityTypeManager()->getStorage($entity_type)->loadMultiple());
+      $expected_count = isset($original_counts[$entity_type]) ? $original_counts[$entity_type] : 0;
+      $this->assertEqual($expected_count, $real_count, "Found $real_count $entity_type entities, expected $expected_count.");
+    }
   }
 
   /**