diff --git a/core/tests/Drupal/Tests/Core/Entity/KeyValueStore/KeyValueEntityStorageTest.php b/core/tests/Drupal/Tests/Core/Entity/KeyValueStore/KeyValueEntityStorageTest.php
index e86b7f2..3b33052 100644
--- a/core/tests/Drupal/Tests/Core/Entity/KeyValueStore/KeyValueEntityStorageTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/KeyValueStore/KeyValueEntityStorageTest.php
@@ -85,9 +85,7 @@ public static function getInfo() {
   protected function setUp() {
     parent::setUp();
 
-    $entity = $this->getMockForAbstractClass('Drupal\Core\Entity\Entity', array(), '', FALSE, TRUE, TRUE, array('onSaveOrDelete'));
-    $entity->expects($this->any())
-      ->method('onSaveOrDelete');
+    $entity = $this->getMockEntity();
 
     $this->entityType = $this->getMock('Drupal\Core\Entity\EntityTypeInterface');
     $this->entityType->expects($this->any())
@@ -406,11 +404,20 @@ public function testDeleteNothing() {
     $this->entityStorage->delete(array());
   }
 
+  /**
+   * Creates an entity with specific methods mocked.
+   *
+   * @param array $methods
+   *   The methods to mock.
+   *
+   * @return \Drupal\Core\Entity\EntityInterface|\PHPUnit_Framework_MockObject_MockObject
+   */
   protected function getMockEntity($methods = array()) {
+    // These methods are not directly related to this storage controller, and
+    // they contain direct calls to \Drupal::entityManager(), so we mock them.
     $methods[] = 'onSaveOrDelete';
+    $methods[] = 'onUpdateBundleEntity';
     $entity = $this->getMockForAbstractClass('Drupal\Core\Entity\Entity', array(), '', FALSE, TRUE, TRUE, $methods);
-    $entity->expects($this->any())
-      ->method('onSaveOrDelete');
     return $entity;
   }