diff --git a/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php b/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
index 9f1b644..b9c5c19 100644
--- a/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
+++ b/src/EntityEmbedDisplay/EntityEmbedDisplayBase.php
@@ -68,6 +68,7 @@ abstract class EntityEmbedDisplayBase extends PluginBase implements ContainerFac
    * {@inheritdoc}
    */
   public function access(AccountInterface $account = NULL) {
+    // @todo Add a hook_entity_embed_display_access()?
     // Check that the plugin's registered entity types matches the current
     // entity type.
     return AccessResult::allowedIf($this->isValidEntityType())
diff --git a/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php b/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php
index e66c5fc..71146d6 100644
--- a/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php
+++ b/src/Plugin/entity_embed/EntityEmbedDisplay/ImageFieldFormatter.php
@@ -90,6 +90,13 @@ class ImageFieldFormatter extends FileFieldFormatter {
   }
 
   protected function isValidImage() {
+    // If entity type is not file we have to return early to prevent fatal in
+    // the condition above. Access should already be forbidden at this point,
+    // which means this won't have any effect.
+    // @see EntityEmbedDisplayBase::access()
+    if ($this->getEntityTypeFromContext() != 'file') {
+      return AccessResult::forbidden();
+    }
     $access = AccessResult::allowed();
 
     // @todo needs cacheability metadata for getEntityFromContext.
diff --git a/src/Tests/EntityEmbedDialogTest.php b/src/Tests/EntityEmbedDialogTest.php
index fcc80cd..0bdfd38 100644
--- a/src/Tests/EntityEmbedDialogTest.php
+++ b/src/Tests/EntityEmbedDialogTest.php
@@ -1,10 +1,5 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Tests\EntityEmbedDialogTest.
- */
-
 namespace Drupal\entity_embed\Tests;
 
 use Drupal\editor\Entity\Editor;
@@ -17,6 +12,13 @@ use Drupal\editor\Entity\Editor;
 class EntityEmbedDialogTest extends EntityEmbedTestBase {
 
   /**
+   * Modules to enable.
+   *
+   * @var array
+   */
+  public static $modules = ['image'];
+
+  /**
    * Tests the entity embed dialog.
    */
   public function testEntityEmbedDialog() {
@@ -102,6 +104,17 @@ class EntityEmbedDialogTest extends EntityEmbedTestBase {
   }
 
   /**
+   * Tests entity embed functionality.
+   */
+  public function testEntityEmbedFunctionality() {
+    $edit = [
+      'entity_id' => $this->node->getTitle() . ' (' . $this->node->id() . ')',
+    ];
+    $this->getEmbedDialog('custom_format', 'node');
+    $this->drupalPostForm(NULL, $edit, t('Next'));
+  }
+
+  /**
    * Retrieves an embed dialog based on given parameters.
    *
    * @param string $filter_format_id