diff --git a/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageSchemaTest.php b/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageSchemaTest.php
index 1e033af..32f0bd9 100644
--- a/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageSchemaTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageSchemaTest.php
@@ -14,6 +14,7 @@
 
 /**
  * @coversDefaultClass \Drupal\Core\Entity\Sql\SqlContentEntityStorageSchema
+ * @group Entity
  */
 class SqlContentEntityStorageSchemaTest extends UnitTestCase {
 
@@ -1083,15 +1084,15 @@ public function providerTestRequiresEntityDataMigration() {
 
     return [
       // Case 1: same storage class, ::hasData() === TRUE.
-      [$updated_entity_type_definition, $original_entity_type_definition, TRUE, TRUE],
+      [$updated_entity_type_definition, $original_entity_type_definition, TRUE, NULL, TRUE],
       // Case 2: same storage class, ::hasData() === FALSE.
-      [$updated_entity_type_definition, $original_entity_type_definition, FALSE, FALSE],
+      [$updated_entity_type_definition, $original_entity_type_definition, FALSE, NULL, FALSE],
       // Case 3: different storage class, original storage class does not exist.
-      [$updated_entity_type_definition, $original_entity_type_definition_other_nonexisting, NULL, TRUE],
+      [$updated_entity_type_definition, $original_entity_type_definition_other_nonexisting, NULL, NULL, TRUE],
       // Case 4: different storage class, original storage class exists, ::hasData() === TRUE.
-      [$updated_entity_type_definition, $original_entity_type_definition_other_existing, TRUE, TRUE],
+      [$updated_entity_type_definition, $original_entity_type_definition_other_existing, NULL, TRUE, TRUE],
       // Case 5: different storage class, original storage class exists, ::hasData() === FALSE.
-      [$updated_entity_type_definition, $original_entity_type_definition_other_existing, FALSE, FALSE],
+      [$updated_entity_type_definition, $original_entity_type_definition_other_existing, NULL, FALSE, FALSE],
     ];
   }
 
@@ -1100,12 +1101,20 @@ public function providerTestRequiresEntityDataMigration() {
    *
    * @dataProvider providerTestRequiresEntityDataMigration
    */
-  public function testRequiresEntityDataMigration($updated_entity_type_definition, $original_entity_type_definition, $storage_has_data, $migration_required) {
+  public function testRequiresEntityDataMigration($updated_entity_type_definition, $original_entity_type_definition, $storage_has_data, $original_storage_has_data, $migration_required) {
     $this->entityType = new ContentEntityType(array(
       'id' => 'entity_test',
       'entity_keys' => array('id' => 'id'),
     ));
 
+    $original_storage = $this->getMockBuilder('Drupal\Core\Entity\Sql\SqlContentEntityStorage')
+      ->disableOriginalConstructor()
+      ->getMock();
+
+    $original_storage->expects($this->exactly(is_null($original_storage_has_data) ? 0 : 1))
+      ->method('hasData')
+      ->willReturn($original_storage_has_data);
+
     // Assert hasData() is not called when $storage_has_data === NULL.
     $this->storage->expects($this->exactly(is_null($storage_has_data) ? 0 : 1))
       ->method('hasData')
@@ -1117,7 +1126,7 @@ public function testRequiresEntityDataMigration($updated_entity_type_definition,
 
     $this->entityManager->expects($this->any())
       ->method('createHandlerInstance')
-      ->willReturn($this->storage);
+      ->willReturn($original_storage);
 
     $this->storageSchema = $this->getMockBuilder('Drupal\Core\Entity\Sql\SqlContentEntityStorageSchema')
       ->setConstructorArgs(array($this->entityManager, $this->entityType, $this->storage, $connection))