diff --git a/core/tests/Drupal/Tests/Core/Entity/EntityManagerTest.php b/core/tests/Drupal/Tests/Core/Entity/EntityManagerTest.php
index 9a67a26..a3f8568 100644
--- a/core/tests/Drupal/Tests/Core/Entity/EntityManagerTest.php
+++ b/core/tests/Drupal/Tests/Core/Entity/EntityManagerTest.php
@@ -1241,19 +1241,15 @@ public function testOnFieldDefinitionCreateNewField() {
     $field_definition->getTargetBundle()->willReturn('test_bundle');
     $field_definition->getName()->willReturn('test_field');
     $field_definition->getType()->willReturn('test_type');
+    $field = $field_definition->reveal();
 
-    $class = $this->getMockClass(DynamicallyFieldableEntityStorageInterface::class);
+    $storage = $this->prophesize(DynamicallyFieldableEntityStorageInterface::class);
+    $storage->onFieldDefinitionCreate($field)
+      ->shouldBeCalledTimes(1);
+    $this->entityManager->setHandler('storage', 'test_entity_type', $storage->reveal());
     $entity = $this->prophesize(EntityTypeInterface::class);
-    $entity->getHandlerClass('storage')->willReturn($class);
     $this->setUpEntityManager(array('test_entity_type' => $entity));
 
-    // The entity manager will instantiate a new object with the given class
-    // name. Define the mock expectations on that.
-    $storage = $this->entityManager->getStorage('test_entity_type');
-    $storage->expects($this->once())
-      ->method('onFieldDefinitionCreate')
-      ->with($field_definition->reveal());
-
     // Set up the stored bundle field map.
     $key_value_store = $this->prophesize(KeyValueStoreInterface::class);
     $this->keyValueFactory->get('entity.definitions.bundle_field_map')->willReturn($key_value_store->reveal());
@@ -1265,7 +1261,7 @@ public function testOnFieldDefinitionCreateNewField() {
         ],
       ])->shouldBeCalled();
 
-    $this->entityManager->onFieldDefinitionCreate($field_definition->reveal());
+    $this->entityManager->onFieldDefinitionCreate($field);
   }
 
   /**
@@ -1276,19 +1272,15 @@ public function testOnFieldDefinitionCreateExistingField() {
     $field_definition->getTargetEntityTypeId()->willReturn('test_entity_type');
     $field_definition->getTargetBundle()->willReturn('test_bundle');
     $field_definition->getName()->willReturn('test_field');
+    $field = $field_definition->reveal();
 
-    $class = $this->getMockClass(DynamicallyFieldableEntityStorageInterface::class);
+    $storage = $this->prophesize(DynamicallyFieldableEntityStorageInterface::class);
+    $storage->onFieldDefinitionCreate($field)
+      ->shouldBeCalledTimes(1);
+    $this->entityManager->setHandler('storage', 'test_entity_type', $storage->reveal());
     $entity = $this->prophesize(EntityTypeInterface::class);
-    $entity->getHandlerClass('storage')->willReturn($class);
     $this->setUpEntityManager(array('test_entity_type' => $entity));
 
-    // The entity manager will instantiate a new object with the given class
-    // name. Define the mock expectations on that.
-    $storage = $this->entityManager->getStorage('test_entity_type');
-    $storage->expects($this->once())
-      ->method('onFieldDefinitionCreate')
-      ->with($field_definition->reveal());
-
     // Set up the stored bundle field map.
     $key_value_store = $this->prophesize(KeyValueStoreInterface::class);
     $this->keyValueFactory->get('entity.definitions.bundle_field_map')->willReturn($key_value_store->reveal());
@@ -1306,7 +1298,7 @@ public function testOnFieldDefinitionCreateExistingField() {
       ])
       ->shouldBeCalled();
 
-    $this->entityManager->onFieldDefinitionCreate($field_definition->reveal());
+    $this->entityManager->onFieldDefinitionCreate($field);
   }
 
   /**
@@ -1315,20 +1307,15 @@ public function testOnFieldDefinitionCreateExistingField() {
   public function testOnFieldDefinitionUpdate() {
     $field_definition = $this->prophesize(FieldDefinitionInterface::class);
     $field_definition->getTargetEntityTypeId()->willReturn('test_entity_type');
+    $field = $field_definition->reveal();
 
-    $class = $this->getMockClass(DynamicallyFieldableEntityStorageInterface::class);
+    $storage = $this->prophesize(DynamicallyFieldableEntityStorageInterface::class);
+    $storage->onFieldDefinitionUpdate($field, $field)
+      ->shouldBeCalledTimes(1);
     $entity = $this->prophesize(EntityTypeInterface::class);
-    $entity->getHandlerClass('storage')->willReturn($class);
+    $this->entityManager->setHandler('storage', 'test_entity_type', $storage->reveal());
     $this->setUpEntityManager(array('test_entity_type' => $entity));
-
-    // The entity manager will instantiate a new object with the given class
-    // name. Define the mock expectations on that.
-    $storage = $this->entityManager->getStorage('test_entity_type');
-    $storage->expects($this->once())
-      ->method('onFieldDefinitionUpdate')
-      ->with($field_definition->reveal());
-
-    $this->entityManager->onFieldDefinitionUpdate($field_definition->reveal(), $field_definition->reveal());
+    $this->entityManager->onFieldDefinitionUpdate($field, $field);
   }
 
   /**
@@ -1339,19 +1326,15 @@ public function testOnFieldDefinitionDeleteMultipleBundles() {
     $field_definition->getTargetEntityTypeId()->willReturn('test_entity_type');
     $field_definition->getTargetBundle()->willReturn('test_bundle');
     $field_definition->getName()->willReturn('test_field');
+    $field = $field_definition->reveal();
 
-    $class = $this->getMockClass(DynamicallyFieldableEntityStorageInterface::class);
+    $storage = $this->prophesize(DynamicallyFieldableEntityStorageInterface::class);
+    $storage->onFieldDefinitionDelete($field)
+      ->shouldBeCalledTimes(1);
+    $this->entityManager->setHandler('storage', 'test_entity_type', $storage->reveal());
     $entity = $this->prophesize(EntityTypeInterface::class);
-    $entity->getHandlerClass('storage')->willReturn($class);
     $this->setUpEntityManager(array('test_entity_type' => $entity));
 
-    // The entity manager will instantiate a new object with the given class
-    // name. Define the mock expectations on that.
-    $storage = $this->entityManager->getStorage('test_entity_type');
-    $storage->expects($this->once())
-      ->method('onFieldDefinitionDelete')
-      ->with($field_definition->reveal());
-
     // Set up the stored bundle field map.
     $key_value_store = $this->prophesize(KeyValueStoreInterface::class);
     $this->keyValueFactory->get('entity.definitions.bundle_field_map')->willReturn($key_value_store->reveal());
@@ -1373,7 +1356,7 @@ public function testOnFieldDefinitionDeleteMultipleBundles() {
       ])
       ->shouldBeCalled();
 
-    $this->entityManager->onFieldDefinitionDelete($field_definition->reveal());
+    $this->entityManager->onFieldDefinitionDelete($field);
   }
 
 
@@ -1385,19 +1368,15 @@ public function testOnFieldDefinitionDeleteSingleBundles() {
     $field_definition->getTargetEntityTypeId()->willReturn('test_entity_type');
     $field_definition->getTargetBundle()->willReturn('test_bundle');
     $field_definition->getName()->willReturn('test_field');
+    $field = $field_definition->reveal();
 
-    $class = $this->getMockClass(DynamicallyFieldableEntityStorageInterface::class);
+    $storage = $this->prophesize(DynamicallyFieldableEntityStorageInterface::class);
+    $storage->onFieldDefinitionDelete($field)
+      ->shouldBeCalledTimes(1);
+    $this->entityManager->setHandler('storage', 'test_entity_type', $storage->reveal());
     $entity = $this->prophesize(EntityTypeInterface::class);
-    $entity->getHandlerClass('storage')->willReturn($class);
     $this->setUpEntityManager(array('test_entity_type' => $entity));
 
-    // The entity manager will instantiate a new object with the given class
-    // name. Define the mock expectations on that.
-    $storage = $this->entityManager->getStorage('test_entity_type');
-    $storage->expects($this->once())
-      ->method('onFieldDefinitionDelete')
-      ->with($field_definition->reveal());
-
     // Set up the stored bundle field map.
     $key_value_store = $this->prophesize(KeyValueStoreInterface::class);
     $this->keyValueFactory->get('entity.definitions.bundle_field_map')->willReturn($key_value_store->reveal());
@@ -1415,7 +1394,7 @@ public function testOnFieldDefinitionDeleteSingleBundles() {
       ])
       ->shouldBeCalled();
 
-    $this->entityManager->onFieldDefinitionDelete($field_definition->reveal());
+    $this->entityManager->onFieldDefinitionDelete($field);
   }
 
   /**
@@ -1570,6 +1549,20 @@ public function setDiscovery(DiscoveryInterface $discovery) {
   }
 
   /**
+   * Allow settings of mock handlers.
+   *
+   * @param string $entity_type
+   *   The entity type for this handler.
+   * @param string $handler_type
+   *   The handler type to create an instance for.
+   * @param object $handler
+   *   A handler instance.
+   */
+  public function setHandler($handler_type, $entity_type, $handler) {
+    $this->handlers[$handler_type][$entity_type] = $handler;
+  }
+
+  /**
    * Allows the static caches to be cleared.
    */
   public function testClearEntityFieldInfo() {