diff --git a/src/EntityEmbedDisplay/EntityEmbedDisplayManager.php b/src/EntityEmbedDisplay/EntityEmbedDisplayManager.php
index 2bfa4088..1c49a25a 100644
--- a/src/EntityEmbedDisplay/EntityEmbedDisplayManager.php
+++ b/src/EntityEmbedDisplay/EntityEmbedDisplayManager.php
@@ -19,7 +19,7 @@
  * @see \Drupal\entity_embed\Annotation\EntityEmbedDisplay
  * @see \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayInterface
  */
-class EntityEmbedDisplayManager extends DefaultPluginManager {
+class EntityEmbedDisplayManager extends DefaultPluginManager implements EntityEmbedDisplayManagerInterface {
 
   /**
    * Constructs a new class instance.
@@ -52,16 +52,7 @@ public function processDefinition(&$definition, $plugin_id) {
   }
 
   /**
-   * Determines plugins whose constraints are satisfied by a set of contexts.
-   *
-   * @param array $contexts
-   *   An array of contexts.
-   *
-   * @return array
-   *   An array of plugin definitions.
-   *
-   * @todo At some point convert this to use ContextAwarePluginManagerTrait
-   * @see https://drupal.org/node/2277981
+   * @{inheritdoc}
    */
   public function getDefinitionsForContexts(array $contexts = array()) {
     $definitions = $this->getDefinitions();
@@ -83,13 +74,7 @@ public function getDefinitionsForContexts(array $contexts = array()) {
   }
 
   /**
-   * Provides a list of plugins that can be used for a certain entity.
-   *
-   * @param \Drupal\Core\Entity\EntityInterface $entity
-   *   An entity object.
-   *
-   * @return array
-   *   An array of valid plugin labels, keyed by plugin ID.
+   * @{inheritdoc}
    */
   public function getDefinitionOptionsForEntity(EntityInterface $entity) {
     $definitions = $this->getDefinitionsForContexts(array('entity' => $entity));
@@ -99,13 +84,7 @@ public function getDefinitionOptionsForEntity(EntityInterface $entity) {
   }
 
   /**
-   * Provides a list of plugins that can be used for a certain entity type.
-   *
-   * @param string $entity_type
-   *   The entity type id.
-   *
-   * @return array
-   *   An array of valid plugin labels, keyed by plugin ID.
+   * @{inheritdoc}
    */
   public function getDefinitionOptionsForEntityType($entity_type) {
     $definitions = $this->getDefinitionsForContexts(array('entity_type' => $entity_type));
diff --git a/src/EntityEmbedDisplay/EntityEmbedDisplayManagerInterface.php b/src/EntityEmbedDisplay/EntityEmbedDisplayManagerInterface.php
new file mode 100644
index 00000000..24a5acaa
--- /dev/null
+++ b/src/EntityEmbedDisplay/EntityEmbedDisplayManagerInterface.php
@@ -0,0 +1,64 @@
+<?php
+
+/**
+ * @file
+ * Contains \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface.
+ */
+
+namespace Drupal\entity_embed\EntityEmbedDisplay;
+
+use Drupal\Component\Plugin\PluginManagerInterface;
+use Drupal\Core\Entity\EntityInterface;
+
+/**
+ * Provides an Entity embed display plugin manager.
+ *
+ * @see \Drupal\entity_embed\Annotation\EntityEmbedDisplay
+ * @see \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayInterface
+ */
+interface EntityEmbedDisplayManagerInterface extends PluginManagerInterface {
+
+  /**
+   * The default display plugin.
+   *
+   * @var string
+   */
+  const DEFAULT_PLUGIN_ID = 'entity_reference:entity_reference_entity_view';
+
+  /**
+   * Determines plugins whose constraints are satisfied by a set of contexts.
+   *
+   * @param array $contexts
+   *   An array of contexts.
+   *
+   * @return array
+   *   An array of plugin definitions.
+   *
+   * @todo At some point convert this to use ContextAwarePluginManagerTrait
+   * @see https://drupal.org/node/2277981
+   */
+  public function getDefinitionsForContexts(array $contexts = array());
+
+  /**
+   * Provides a list of plugins that can be used for a certain entity.
+   *
+   * @param \Drupal\Core\Entity\EntityInterface $entity
+   *   An entity object.
+   *
+   * @return array
+   *   An array of valid plugin labels, keyed by plugin ID.
+   */
+  public function getDefinitionOptionsForEntity(EntityInterface $entity);
+
+  /**
+   * Provides a list of plugins that can be used for a certain entity type.
+   *
+   * @param string $entity_type
+   *   The entity type id.
+   *
+   * @return array
+   *   An array of valid plugin labels, keyed by plugin ID.
+   */
+  public function getDefinitionOptionsForEntityType($entity_type);
+
+}
diff --git a/src/EntityHelperTrait.php b/src/EntityHelperTrait.php
index 13ac8802..2f68521b 100644
--- a/src/EntityHelperTrait.php
+++ b/src/EntityHelperTrait.php
@@ -12,7 +12,7 @@
 use Drupal\Core\Entity\EntityStorageException;
 use Drupal\Core\Extension\ModuleHandlerInterface;
 use Drupal\Core\Render\RendererInterface;
-use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
+use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface;
 
 /**
  * Wrapper methods for entity loading and rendering.
@@ -41,7 +41,7 @@ trait EntityHelperTrait {
   /**
    * The display plugin manager.
    *
-   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager.
+   * @var \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface.
    */
   protected $displayPluginManager;
 
@@ -162,25 +162,7 @@ protected function renderEntity(EntityInterface $entity, $view_mode, $langcode =
    *   The HTML of the entity rendered with the display plugin.
    */
   protected function renderEntityEmbed(EntityInterface $entity, array $context = array()) {
-    // Support the deprecated view-mode data attribute.
-    if (isset($context['data-view-mode']) && !isset($context['data-entity-embed-display']) && !isset($context['data-entity-embed-settings'])) {
-      $context['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
-      $context['data-entity-embed-settings'] = ['view_mode' => &$context['data-view-mode']];
-    }
-
-    // Merge in default attributes.
-    $context += array(
-      'data-entity-id' => $entity->id(),
-      'data-entity-type' => $entity->getEntityTypeId(),
-      'data-entity-embed-display' => 'entity_reference:entity_reference_entity_view',
-      'data-entity-embed-settings' => array(),
-    );
-
-    // The default display plugin has been deprecated by the rendered entity
-    // field formatter.
-    if ($context['data-entity-embed-display'] === 'default') {
-      $context['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
-    }
+    $this->prepareEmbedContext($context, $entity);
 
     // Allow modules to alter the entity prior to embed rendering.
     $this->moduleHandler()->alter(array("{$context['data-entity-type']}_embed_context", 'entity_embed_context'), $context, $entity);
@@ -200,6 +182,27 @@ protected function renderEntityEmbed(EntityInterface $entity, array $context = a
     return $entity_output;
   }
 
+  protected function prepareEmbedContext(array &$context, EntityInterface $entity = null) {
+    // Merge in default attributes.
+    $context += array(
+      'data-entity-embed-display' => EntityEmbedDisplayManagerInterface::DEFAULT_PLUGIN_ID,
+      'data-entity-embed-settings' => array(),
+      'data-entity-id' => $entity ? $entity->id() : '',
+      'data-entity-uuid' => $entity ? $entity->uuid() : '',
+      'data-entity-type' => $entity ? $entity->getEntityTypeId() : '',
+    );
+
+    if ($context['data-entity-embed-display'] === 'default') {
+      // The default display plugin has been deprecated by the rendered entity
+      // field formatter.
+      $context['data-entity-embed-display'] = EntityEmbedDisplayManagerInterface::DEFAULT_PLUGIN_ID;
+    }
+    elseif ($context['data-entity-embed-display'] === EntityEmbedDisplayManagerInterface::DEFAULT_PLUGIN_ID && isset($context['data-view-mode']) && empty($context['data-entity-embed-settings'])) {
+      // Support the deprecated view-mode data attribute.
+      $context['data-entity-embed-settings']['view_mode'] = &$context['data-view-mode'];
+    }
+  }
+
   /**
    * Renders an entity using an EntityEmbedDisplay plugin.
    *
@@ -287,7 +290,7 @@ public function setModuleHandler(ModuleHandlerInterface $module_handler) {
   /**
    * Returns the display plugin manager.
    *
-   * @return \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager
+   * @return \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface
    *   The display plugin manager.
    */
   protected function displayPluginManager() {
@@ -300,12 +303,12 @@ protected function displayPluginManager() {
   /**
    * Sets the display plugin manager service.
    *
-   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $display_plugin_manager
+   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface $display_plugin_manager
    *   The display plugin manager service.
    *
    * @return self
    */
-  public function setDisplayPluginManager(EntityEmbedDisplayManager $display_plugin_manager) {
+  public function setDisplayPluginManager(EntityEmbedDisplayManagerInterface $display_plugin_manager) {
     $this->displayPluginManager = $display_plugin_manager;
     return $this;
   }
diff --git a/src/Form/EntityEmbedDialog.php b/src/Form/EntityEmbedDialog.php
index 9b2f5e4b..12672243 100644
--- a/src/Form/EntityEmbedDialog.php
+++ b/src/Form/EntityEmbedDialog.php
@@ -17,7 +17,7 @@
 use Drupal\editor\Ajax\EditorDialogSave;
 use Drupal\editor\EditorInterface;
 use Drupal\embed\EmbedButtonInterface;
-use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager;
+use Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface;
 use Drupal\entity_embed\EntityHelperTrait;
 use Drupal\Component\Serialization\Json;
 use Symfony\Component\DependencyInjection\ContainerInterface;
@@ -38,12 +38,12 @@ class EntityEmbedDialog extends FormBase {
   /**
    * Constructs a EntityEmbedDialog object.
    *
-   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManager $plugin_manager
-   *   The Module Handler.
+   * @param \Drupal\entity_embed\EntityEmbedDisplay\EntityEmbedDisplayManagerInterface $plugin_manager
+   *   The display plugin manager.
    * @param \Drupal\Core\Form\FormBuilderInterface $form_builder
-   *   The Form Builder.
+   *   The form builder.
    */
-  public function __construct(EntityEmbedDisplayManager $plugin_manager, FormBuilderInterface $form_builder) {
+  public function __construct(EntityEmbedDisplayManagerInterface $plugin_manager, FormBuilderInterface $form_builder) {
     $this->setDisplayPluginManager($plugin_manager);
     $this->formBuilder = $form_builder;
   }
@@ -76,10 +76,7 @@ public function getFormId() {
   public function buildForm(array $form, FormStateInterface $form_state, EditorInterface $editor = NULL, EmbedButtonInterface $embed_button = NULL) {
     $values = $form_state->getValues();
     $input = $form_state->getUserInput();
-    // Set embed button element in form state, so that it can be used later in
-    // validateForm() function.
-    $form_state->set('embed_button', $embed_button);
-    $form_state->set('editor', $editor);
+
     // Initialize entity element with form attributes, if present.
     $entity_element = empty($values['attributes']) ? array() : $values['attributes'];
     // The default values are set directly from \Drupal::request()->request,
@@ -90,14 +87,14 @@ public function buildForm(array $form, FormStateInterface $form_state, EditorInt
     $entity_element += $form_state->get('entity_element');
     $entity_element += array(
       'data-entity-type' => $embed_button->getTypeSetting('entity_type'),
-      'data-entity-uuid' => '',
-      'data-entity-id' => '',
-      'data-entity-embed-display' => 'entity_reference:entity_reference_entity_view',
-      'data-entity-embed-settings' => array(),
       'data-align' => '',
     );
-    $form_state->set('entity_element', $entity_element);
+    $this->prepareEmbedContext($entity_element, $form_state->get('entity'));
+
+    $form_state->set('embed_button', $embed_button);
+    $form_state->set('editor', $editor);
     $form_state->set('entity', $this->loadEntity($entity_element['data-entity-type'], $entity_element['data-entity-uuid'] ?: $entity_element['data-entity-id']));
+    $form_state->set('entity_element', $entity_element);
 
     if (!$form_state->get('step')) {
       // If an entity has been selected, then always skip to the embed options.
@@ -250,12 +247,6 @@ public function buildEmbedStep(array $form, FormStateInterface $form_state) {
       $entity_element['data-entity-embed-display'] = key($display_plugin_options);
     }
 
-    // The default display plugin has been deprecated by the rendered entity
-    // field formatter.
-    if ($entity_element['data-entity-embed-display'] === 'default') {
-      $entity_element['data-entity-embed-display'] = 'entity_reference:entity_reference_entity_view';
-    }
-
     $form['attributes']['data-entity-embed-display'] = array(
       '#type' => 'select',
       '#title' => $this->t('Display as'),