diff --git a/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php b/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
index 50f4411..3c55f0b 100644
--- a/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
+++ b/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
@@ -11,7 +11,7 @@ use Drupal\Core\Language\LanguageManagerInterface;
 use Drupal\Core\Plugin\ContainerFactoryPluginInterface;
 use Drupal\Core\Plugin\PluginBase;
 use Drupal\Core\Session\AccountInterface;
-use Drupal\entity_embed\EntityHelperTrait;
+use Drupal\entity_embed\EntityEmbedHelperTrait;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
@@ -25,7 +25,7 @@ use Symfony\Component\DependencyInjection\ContainerInterface;
   * @ingroup entity_embed_api
  */
 abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFactoryPluginInterface, EntityEmbedDisplayInterface {
-  use EntityHelperTrait;
+  use EntityEmbedHelperTrait;
 
   /**
    * The entity type manager service.
diff --git a/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php b/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php
index 3920b88..595fbda 100644
--- a/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php
+++ b/src/EntityEmbedDisplay/FieldFormatterEntityEmbedDisplayBase.php
@@ -18,13 +18,6 @@ abstract class FieldFormatterEntityEmbedDisplayBase extends EntityEmbedDisplayBa
   use PluginDependencyTrait;
 
   /**
-   * The entity type manager service.
-   *
-   * @var \Drupal\Core\Entity\EntityTypeManagerInterface
-   */
-  protected $entityTypeManager;
-
-  /**
    * The field formatter plugin manager.
    *
    * @var \Drupal\Core\Field\FormatterPluginManager
@@ -67,7 +60,6 @@ abstract class FieldFormatterEntityEmbedDisplayBase extends EntityEmbedDisplayBa
   public function __construct(array $configuration, $plugin_id, $plugin_definition, EntityTypeManagerInterface $entity_type_manager, FormatterPluginManager $formatter_plugin_manager, TypedDataManager $typed_data_manager, LanguageManagerInterface $language_manager) {
     $this->formatterPluginManager = $formatter_plugin_manager;
     $this->setConfiguration($configuration);
-    $this->entityTypeManager = $entity_type_manager;
     $this->typedDataManager = $typed_data_manager;
     parent::__construct($configuration, $plugin_id, $plugin_definition, $entity_type_manager, $language_manager);
   }
diff --git a/src/EntityHelperTrait.php b/src/EntityEmbedHelperTrait.php
similarity index 64%
rename from src/EntityHelperTrait.php
rename to src/EntityEmbedHelperTrait.php
index 3cc76d9..37d63bb 100644
--- a/src/EntityHelperTrait.php
+++ b/src/EntityEmbedHelperTrait.php
@@ -4,13 +4,81 @@ namespace Drupal\entity_embed;
 
 use Drupal\Component\Utility\Html;
 use Drupal\Core\Entity\EntityInterface;
+use Drupal\Core\Extension\ModuleHandlerInterface;
+use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
 
 /**
  * Wrapper methods for entity rendering.
  *
  * @internal
  */
-trait EntityHelperTrait {
+trait EntityEmbedHelperTrait {
+
+  /**
+   * Teh Entity Embed Display plugin manager.
+   *
+   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
+   */
+  protected $displayPluginManager;
+
+  /**
+   * The module handler service.
+   *
+   * @var \Drupal\Core\Extension\ModuleHandlerInterface.
+   */
+  protected $moduleHandler;
+
+  /**
+   * Returns the Entity Embed Display plugin manager.
+   *
+   * @return \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
+   *   The Entity Embed Display plugin manager.
+   */
+  protected function displayPluginManager() {
+    if (!isset($this->displayPluginManager)) {
+      $this->displayPluginManager = \Drupal::service('plugin.manager.entity_embed.display');
+    }
+    return $this->displayPluginManager;
+  }
+
+  /**
+   * Sets the Entity Embed Display plugin manager.
+   *
+   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $display_plugin_manager
+   *   The Entity Embed Display plugin manager.
+   *
+   * @return self
+   */
+  public function setDisplayPluginManager(EntityEmbedDisplayManager $display_plugin_manager) {
+    $this->displayPluginManager = $display_plugin_manager;
+    return $this;
+  }
+
+  /**
+   * Returns the module handler.
+   *
+   * @return \Drupal\Core\Extension\ModuleHandlerInterface
+   *   The module handler.
+   */
+  protected function moduleHandler() {
+    if (!isset($this->moduleHandler)) {
+      $this->moduleHandler = \Drupal::moduleHandler();
+    }
+    return $this->moduleHandler;
+  }
+
+  /**
+   * Sets the module handler service.
+   *
+   * @param \Drupal\Core\Extension\ModuleHandlerInterface $module_handler
+   *   The module handler service.
+   *
+   * @return self
+   */
+  public function setModuleHandler(ModuleHandlerInterface $module_handler) {
+    $this->moduleHandler = $module_handler;
+    return $this;
+  }
 
   /**
    * Builds the render array for an embedded entity.
@@ -53,7 +121,7 @@ trait EntityHelperTrait {
     }
 
     // Allow modules to alter the entity prior to embed rendering.
-    \Drupal::moduleHandler()->alter(array("{$context['data-entity-type']}_embed_context", 'entity_embed_context'), $context, $entity);
+    $this->moduleHandler()->alter(array("{$context['data-entity-type']}_embed_context", 'entity_embed_context'), $context, $entity);
 
     // Build and render the Entity Embed Display plugin, allowing modules to
     // alter the result before rendering.
@@ -87,7 +155,7 @@ trait EntityHelperTrait {
     $build['#access'] = $entity->access('view', NULL, TRUE);
 
     // @todo Should this hook get invoked if $build is an empty array?
-    \Drupal::moduleHandler()->alter(array("{$context['data-entity-type']}_embed", 'entity_embed'), $build, $entity, $context);
+    $this->moduleHandler()->alter(array("{$context['data-entity-type']}_embed", 'entity_embed'), $build, $entity, $context);
     return $build;
   }
 
@@ -110,7 +178,7 @@ trait EntityHelperTrait {
   protected function buildEntityEmbedDisplayPlugin(EntityInterface $entity, $plugin_id, array $plugin_configuration = array(), array $context = array()) {
     // Build the Entity Embed Display plugin.
     /** @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayBase $display */
-    $display = \Drupal::service('plugin.manager.entity_embed.display')->createInstance($plugin_id, $plugin_configuration);
+    $display = $this->displayPluginManager()->createInstance($plugin_id, $plugin_configuration);
     $display->setContextValue('entity', $entity);
     $display->setAttributes($context);
 
diff --git a/src/Form/EntityEmbedDialog.php b/src/Form/EntityEmbedDialog.php
index 564fc10..ff91751 100644
--- a/src/Form/EntityEmbedDialog.php
+++ b/src/Form/EntityEmbedDialog.php
@@ -21,7 +21,6 @@ use Drupal\embed\EmbedButtonInterface;
 use Drupal\entity_browser\Events\Events;
 use Drupal\entity_browser\Events\RegisterJSCallbacks;
 use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
-use Drupal\entity_embed\EntityHelperTrait;
 use Drupal\Component\Serialization\Json;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 use Symfony\Component\EventDispatcher\EventDispatcherInterface;
@@ -30,7 +29,6 @@ use Symfony\Component\EventDispatcher\EventDispatcherInterface;
  * Provides a form to embed entities by specifying data attributes.
  */
 class EntityEmbedDialog extends FormBase {
-  use EntityHelperTrait;
 
   /**
    * The entity embed display manager.
@@ -581,7 +579,7 @@ class EntityEmbedDialog extends FormBase {
     $entity_element = $form_state->getValue('attributes');
     $entity = $this->entityTypeManager->getStorage($entity_element['data-entity-type'])
       ->loadByProperties(['uuid' => $entity_element['data-entity-uuid']]);
-    $entity = current($entity);
+    $entity = current($entity) ?: NULL;
     $plugin_id = $entity_element['data-entity-embed-display'];
     $plugin_settings = $entity_element['data-entity-embed-settings'] ?: array();
     $display = $this->entityEmbedDisplayManager->createInstance($plugin_id, $plugin_settings);
diff --git a/src/Plugin/Filter/EntityEmbedFilter.php b/src/Plugin/Filter/EntityEmbedFilter.php
index 4ce4829..7dac555 100644
--- a/src/Plugin/Filter/EntityEmbedFilter.php
+++ b/src/Plugin/Filter/EntityEmbedFilter.php
@@ -9,7 +9,7 @@ use Drupal\Core\Plugin\ContainerFactoryPluginInterface;
 use Drupal\Core\Render\BubbleableMetadata;
 use Drupal\Core\Render\RenderContext;
 use Drupal\Core\Render\RendererInterface;
-use Drupal\entity_embed\EntityHelperTrait;
+use Drupal\entity_embed\EntityEmbedHelperTrait;
 use Drupal\entity_embed\Exception\EntityNotFoundException;
 use Drupal\entity_embed\Exception\RecursiveRenderingException;
 use Drupal\filter\FilterProcessResult;
@@ -28,7 +28,7 @@ use Drupal\embed\DomHelperTrait;
  * )
  */
 class EntityEmbedFilter extends FilterBase implements ContainerFactoryPluginInterface {
-  use EntityHelperTrait;
+  use EntityEmbedHelperTrait;
   use DomHelperTrait;
 
   /**
diff --git a/src/Twig/EntityEmbedTwigExtension.php b/src/Twig/EntityEmbedTwigExtension.php
index b836d02..16a0988 100644
--- a/src/Twig/EntityEmbedTwigExtension.php
+++ b/src/Twig/EntityEmbedTwigExtension.php
@@ -4,14 +4,14 @@ namespace Drupal\entity_embed\Twig;
 
 use Drupal\Core\DependencyInjection\ContainerInjectionInterface;
 use Drupal\Core\Entity\EntityTypeManagerInterface;
-use Drupal\entity_embed\EntityHelperTrait;
+use Drupal\entity_embed\EntityEmbedHelperTrait;
 use Symfony\Component\DependencyInjection\ContainerInterface;
 
 /**
  * Provide entity embedding function within Twig templates.
  */
 class EntityEmbedTwigExtension extends \Twig_Extension implements ContainerInjectionInterface {
-  use EntityHelperTrait;
+  use EntityEmbedHelperTrait;
 
   /**
    * The entity type manager service.