diff --git a/core/tests/Drupal/Tests/Core/Entity/EntityTypeTest.php b/core/tests/Drupal/Tests/Core/Entity/EntityTypeTest.php
index b23c796..d633afc 100644
--- a/core/tests/Drupal/Tests/Core/Entity/EntityTypeTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/EntityTypeTest.php
@@ -171,8 +171,7 @@ public function testGetStorageClass() {
   public function testSetStorageClass() {
     $controller = $this->getTestHandlerClass();
     $entity_type = $this->setUpEntityType(array());
-    $entity_type->setStorageClass($controller);
-    $this->assertSame($controller, $entity_type->getHandlerClass('storage'));
+    $this->assertSame($entity_type, $entity_type->setStorageClass($controller));
   }
 
   /**