diff -u b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php
--- b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php
@@ -39,17 +39,7 @@
     $component = \Drupal::service('plugin.manager.field.widget')
       ->prepareConfiguration($field_type, []);
     /** @var \Drupal\Core\Entity\Display\EntityFormDisplayInterface $entity_form_display */
-    $entity_form_display = \Drupal::entityTypeManager()
-      ->getStorage('entity_form_display')
-      ->load('media.' . $media_type_id . '.default');
-    if (!$entity_form_display) {
-      $entity_form_display = EntityFormDisplay::create([
-        'targetEntityType' => 'media',
-        'bundle' => $media_type_id,
-        'mode' => 'default',
-        'status' => TRUE,
-      ]);
-    }
+    $entity_form_display = entity_get_form_display('media', $media_type_id, 'default');
     $entity_form_display->setComponent($field_name, $component)
       ->save();
   }
@@ -78,10 +68,10 @@
    */
   protected function hideMediaTypeFieldWidget($field_name, $media_type_id) {
     /** @var \Drupal\Core\Entity\Display\EntityFormDisplayInterface $entity_form_display */
-    $entity_form_display = \Drupal::entityTypeManager()
-      ->getStorage('entity_form_display')
-      ->load('media.' . $media_type_id . '.default');
-    $entity_form_display->removeComponent($field_name)->save();
+    $entity_form_display = entity_get_form_display('media', $media_type_id, 'default');
+    if ($entity_form_display->getComponent($field_name)) {
+      $entity_form_display->removeComponent($field_name)->save();
+    }
   }
 
   /**