diff --git a/src/Tests/ImageFieldFormatterTest.php b/src/Tests/ImageFieldFormatterTest.php
index de2f6ff..96b86e7 100644
--- a/src/Tests/ImageFieldFormatterTest.php
+++ b/src/Tests/ImageFieldFormatterTest.php
@@ -1,10 +1,5 @@
 <?php
 
-/**
- * @file
- * Contains \Drupal\entity_embed\Tests\ImageFieldFormatterTest.
- */
-
 namespace Drupal\entity_embed\Tests;
 
 use Drupal\Component\Serialization\Json;
@@ -62,7 +57,8 @@ class ImageFieldFormatterTest extends EntityEmbedTestBase {
     // Ensure that correct form attributes are returned for the image plugin.
     $form = array();
     $form_state = new FormState();
-    $display = $this->displayPluginManager()->createInstance('image:image', array());
+    $display = $this->container->get('plugin.manager.entity_embed.display')
+      ->createInstance('image:image', []);
     $display->setContextValue('entity', $this->image);
     $conf_form = $display->buildConfigurationForm($form, $form_state);
     $this->assertIdentical(array_keys($conf_form), array(
diff --git a/src/Tests/ViewModeFieldFormatterTest.php b/src/Tests/ViewModeFieldFormatterTest.php
index 0891901..7712da9 100644
--- a/src/Tests/ViewModeFieldFormatterTest.php
+++ b/src/Tests/ViewModeFieldFormatterTest.php
@@ -27,7 +27,8 @@ class ViewModeFieldFormatterTest extends EntityEmbedTestBase {
     foreach ($this->plugins as $plugin) {
       $form = [];
       $form_state = new FormState();
-      $display = $this->displayPluginManager()->createInstance($plugin, []);
+      $display = $this->container->get('plugin.manager.entity_embed.display')
+        ->createInstance($plugin, []);
       $display->setContextValue('entity', $this->node);
       $conf_form = $display->buildConfigurationForm($form, $form_state);
       $this->assertIdentical(array_keys($conf_form), []);