diff --git a/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageTest.php b/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageTest.php
index 683f9c4..57aa2c3 100644
--- a/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/Sql/SqlContentEntityStorageTest.php
@@ -8,10 +8,10 @@
 namespace Drupal\Tests\Core\Entity\Sql;
 
 use Drupal\Core\Cache\CacheBackendInterface;
-use Drupal\Core\Entity\Sql\SqlContentEntityStorage;
 use Drupal\Core\Entity\EntityInterface;
 use Drupal\Core\Entity\EntityStorageInterface;
-use Drupal\Core\Entity\Schema\SqlContentEntityStorageSchema;
+use Drupal\Core\Entity\Sql\SqlContentEntityStorage;
+use Drupal\Core\Entity\Sql\SqlContentEntityStorageSchema;
 use Drupal\Core\Field\BaseFieldDefinition;
 use Drupal\Core\Language\Language;
 use Drupal\Tests\UnitTestCase;
@@ -336,7 +336,7 @@ public function testOnEntityTypeCreate() {
       ->method('schema')
       ->will($this->returnValue($schema_handler));
 
-    $storage = $this->getMockBuilder('Drupal\Core\Entity\ContentEntityDatabaseStorage')
+    $storage = $this->getMockBuilder('Drupal\Core\Entity\Sql\SqlContentEntityStorage')
       ->setConstructorArgs(array($this->entityType, $this->connection, $this->entityManager, $this->cache))
       ->setMethods(array('schemaHandler'))
       ->getMock();