diff --git a/js/plugins/drupalentity/plugin.js b/js/plugins/drupalentity/plugin.js
index 895a75b..3a95c24 100644
--- a/js/plugins/drupalentity/plugin.js
+++ b/js/plugins/drupalentity/plugin.js
@@ -71,7 +71,7 @@
       // drupallink has a hardcoded integration with drupalimage. Work around that, to reuse the same integration.
 
       var originalGetFocusedWidget = null;
-      if (CKEDITOR.plugins.drupalimage) {
+      if (originalGetFocusedWidget) {
         originalGetFocusedWidget = CKEDITOR.plugins.drupalimage.getFocusedWidget;
       }
       else {
diff --git a/tests/src/FunctionalJavascript/MediaImageTest.php b/tests/src/FunctionalJavascript/MediaImageTest.php
index d619080..eae0455 100644
--- a/tests/src/FunctionalJavascript/MediaImageTest.php
+++ b/tests/src/FunctionalJavascript/MediaImageTest.php
@@ -2,6 +2,7 @@
 
 namespace Drupal\Tests\entity_embed\FunctionalJavascript;
 
+use Drupal\editor\Entity\Editor;
 use Drupal\Component\Utility\Html;
 use Drupal\editor\Entity\Editor;
 use Drupal\entity_embed\Plugin\entity_embed\EntityEmbedDisplay\MediaImageDecorator;