diff --git a/core/modules/media/tests/src/Kernel/MediaKernelTestBase.php b/core/modules/media/tests/src/Kernel/MediaKernelTestBase.php
index d0f62af..5d75bed 100644
--- a/core/modules/media/tests/src/Kernel/MediaKernelTestBase.php
+++ b/core/modules/media/tests/src/Kernel/MediaKernelTestBase.php
@@ -6,6 +6,7 @@
 use Drupal\KernelTests\KernelTestBase;
 use Drupal\media\Entity\Media;
 use Drupal\media\Entity\MediaType;
+use Drupal\media\MediaTypeInterface;
 use org\bovigo\vfs\vfsStream;
 
 /**
@@ -96,14 +97,13 @@ protected function createMediaType($media_source_name) {
    *
    * @param string $filename
    *   String filename with extension.
-   * @param \Drupal\media\Entity\MediaType $media_type
+   * @param \Drupal\media\MediaTypeInterface $media_type
    *   The the media type.
    *
    * @return \Drupal\media\Entity\Media
    *   A media entity.
    */
-  protected function generateMedia($filename, MediaType $media_type) {
-
+  protected function generateMedia($filename, MediaTypeInterface $media_type) {
     vfsStream::setup('drupal_root');
     vfsStream::create([
       'sites' => [