diff --git a/core/modules/media/tests/src/FunctionalJavascript/MediaSourceFileTest.php b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceFileTest.php
index 8c80b83..ac0b339 100644
--- a/core/modules/media/tests/src/FunctionalJavascript/MediaSourceFileTest.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceFileTest.php
@@ -46,7 +46,8 @@ public function testMediaFileSource() {
     // Create a media item.
     $this->drupalGet("media/add/{$media_type_id}");
     $page->attachFileToField("files[{$source_field_id}_0]", \Drupal::service('file_system')->realpath($test_filepath));
-    $assert_session->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForButton('Remove');
+    $this->assertNotEmpty($result);
     $page->pressButton('Save and publish');
 
     $assert_session->addressEquals('media/1');
@@ -63,7 +64,8 @@ public function testMediaFileSource() {
     file_unmanaged_copy($icon_base . '/generic.png', $icon_base . '/text--plain.png');
     $this->drupalGet("media/add/{$media_type_id}");
     $page->attachFileToField("files[{$source_field_id}_0]", \Drupal::service('file_system')->realpath($test_filepath));
-    $assert_session->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForButton('Remove');
+    $this->assertNotEmpty($result);
     $page->pressButton('Save and publish');
     $assert_session->elementAttributeContains('css', '.image-style-thumbnail', 'src', 'text--plain.png');
   }
diff --git a/core/modules/media/tests/src/FunctionalJavascript/MediaSourceImageTest.php b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceImageTest.php
index 0505d3b..be979e7 100644
--- a/core/modules/media/tests/src/FunctionalJavascript/MediaSourceImageTest.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceImageTest.php
@@ -60,7 +60,8 @@ public function testMediaImageSource() {
     // Create a media item.
     $this->drupalGet("media/add/{$media_type_id}");
     $page->attachFileToField("files[{$source_field_id}_0]", \Drupal::root() . '/core/modules/media/tests/fixtures/example_1.jpeg');
-    $assert_session->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForButton('Remove');
+    $this->assertNotEmpty($result);
     $page->fillField("{$source_field_id}[0][alt]", 'Image Alt Text 1');
     $page->pressButton('Save and publish');
 
diff --git a/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php
index 24cd134..c6ea4ec 100644
--- a/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaSourceTestBase.php
@@ -112,11 +112,13 @@ public function doTestCreateMediaType($media_type_id, $source_id, array $provide
     $assert_session->fieldExists('Media source');
     $assert_session->optionExists('Media source', $source_id);
     $page->selectFieldOption('Media source', $source_id);
-    $assert_session->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForElementVisible('css', 'fieldset[data-drupal-selector="edit-source-configuration"]');
+    $this->assertNotEmpty($result);
 
     // Make sure the provided fields are visible on the form.
     foreach ($provided_fields as $provided_field) {
-      $assert_session->selectExists("field_map[{$provided_field}]");
+      $result = $assert_session->waitForElementVisible('css', 'select[name="field_map[' . $provided_field . ']"]');
+      $this->assertNotEmpty($result);
     }
 
     // Save the form to create the type.
diff --git a/core/modules/media/tests/src/FunctionalJavascript/MediaTypeCreationTest.php b/core/modules/media/tests/src/FunctionalJavascript/MediaTypeCreationTest.php
index bd263b5..efbe288 100644
--- a/core/modules/media/tests/src/FunctionalJavascript/MediaTypeCreationTest.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaTypeCreationTest.php
@@ -13,24 +13,27 @@ class MediaTypeCreationTest extends MediaJavascriptTestBase {
    * Tests the media type creation form.
    */
   public function testMediaTypeCreationFormWithDefaultField() {
+    $session = $this->getSession();
+    $page = $session->getPage();
+    $assert_session = $this->assertSession();
+
     $label = 'Type with Default Field';
     $mediaTypeMachineName = str_replace(' ', '_', strtolower($label));
 
     $this->drupalGet('admin/structure/media/add');
-    $page = $this->getSession()->getPage();
 
     // Fill in a label to the media type.
     $page->fillField('label', $label);
     // Wait for machine name generation. Default: waitUntilVisible(), does not
     // work properly.
-    $this->getSession()
-      ->wait(5000, "jQuery('.machine-name-value').text() === '{$mediaTypeMachineName}'");
+    $session->wait(5000, "jQuery('.machine-name-value').text() === '{$mediaTypeMachineName}'");
 
     // Select the media source used by our media type.
-    $this->assertSession()->fieldExists('Media source');
-    $this->assertSession()->optionExists('Media source', 'test');
+    $assert_session->fieldExists('Media source');
+    $assert_session->optionExists('Media source', 'test');
     $page->selectFieldOption('Media source', 'test');
-    $this->assertSession()->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForElementVisible('css', 'fieldset[data-drupal-selector="edit-source-configuration"]');
+    $this->assertNotEmpty($result);
 
     $page->pressButton('Save');
 
@@ -38,51 +41,51 @@ public function testMediaTypeCreationFormWithDefaultField() {
     $this->drupalGet("admin/structure/media/manage/{$mediaTypeMachineName}/fields");
 
     // Check 2nd column of first data row, to be machine name for field name.
-    $this->assertSession()
-      ->elementContains('xpath', '(//table[@id="field-overview"]//tr)[2]//td[2]', 'field_media_test');
+    $assert_session->elementContains('xpath', '(//table[@id="field-overview"]//tr)[2]//td[2]', 'field_media_test');
     // Check 3rd column of first data row, to be correct field type.
-    $this->assertSession()
-      ->elementTextContains('xpath', '(//table[@id="field-overview"]//tr)[2]//td[3]', 'Text (plain)');
+    $assert_session->elementTextContains('xpath', '(//table[@id="field-overview"]//tr)[2]//td[3]', 'Text (plain)');
 
     // Check that the source field is correctly assigned to media type.
     $this->drupalGet("admin/structure/media/manage/{$mediaTypeMachineName}");
 
-    $this->assertSession()->pageTextContains('Test source field is used to store the essential information about the media item.');
+    $assert_session->pageTextContains('Test source field is used to store the essential information about the media item.');
   }
 
   /**
    * Test creation of media type, reusing an existing source field.
    */
   public function testMediaTypeCreationReuseSourceField() {
+    $session = $this->getSession();
+    $page = $session->getPage();
+    $assert_session = $this->assertSession();
+
     // Create a new media type, which should create a new field we can reuse.
     $this->drupalGet('/admin/structure/media/add');
-    $page = $this->getSession()->getPage();
     $page->fillField('label', 'Pastafazoul');
-    $this->getSession()
-      ->wait(5000, "jQuery('.machine-name-value').text() === 'pastafazoul'");
+    $session->wait(5000, "jQuery('.machine-name-value').text() === 'pastafazoul'");
     $page->selectFieldOption('Media source', 'test');
-    $this->assertSession()->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForElementVisible('css', 'fieldset[data-drupal-selector="edit-source-configuration"]');
+    $this->assertNotEmpty($result);
     $page->pressButton('Save');
 
     $label = 'Type reusing Default Field';
     $mediaTypeMachineName = str_replace(' ', '_', strtolower($label));
 
     $this->drupalGet('admin/structure/media/add');
-    $page = $this->getSession()->getPage();
 
     // Fill in a label to the media type.
     $page->fillField('label', $label);
 
     // Wait for machine name generation. Default: waitUntilVisible(), does not
     // work properly.
-    $this->getSession()
-      ->wait(5000, "jQuery('.machine-name-value').text() === '{$mediaTypeMachineName}'");
+    $session->wait(5000, "jQuery('.machine-name-value').text() === '{$mediaTypeMachineName}'");
 
     // Select the media source used by our media type.
-    $this->assertSession()->fieldExists('Media source');
-    $this->assertSession()->optionExists('Media source', 'test');
+    $assert_session->fieldExists('Media source');
+    $assert_session->optionExists('Media source', 'test');
     $page->selectFieldOption('Media source', 'test');
-    $this->assertSession()->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForElementVisible('css', 'fieldset[data-drupal-selector="edit-source-configuration"]');
+    $this->assertNotEmpty($result);
     // Select the existing field for re-use.
     $page->selectFieldOption('source_configuration[source_field]', 'field_media_test');
     $page->pressButton('Save');
@@ -90,9 +93,9 @@ public function testMediaTypeCreationReuseSourceField() {
     // Check that no new fields were created.
     $this->drupalGet("admin/structure/media/manage/{$mediaTypeMachineName}/fields");
     // The reused field should be present...
-    $this->assertSession()->pageTextContains('field_media_test');
+    $assert_session->pageTextContains('field_media_test');
     // ...not a new, unique one.
-    $this->assertSession()->pageTextNotContains('field_media_test_1');
+    $assert_session->pageTextNotContains('field_media_test_1');
   }
 
 }
diff --git a/core/modules/media/tests/src/FunctionalJavascript/MediaUiJavascriptTest.php b/core/modules/media/tests/src/FunctionalJavascript/MediaUiJavascriptTest.php
index b150f89..3655f0b 100644
--- a/core/modules/media/tests/src/FunctionalJavascript/MediaUiJavascriptTest.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaUiJavascriptTest.php
@@ -120,7 +120,8 @@ public function testMediaTypes() {
 
     // Try to change media type and check if new configuration sub-form appears.
     $page->selectFieldOption('source', 'test');
-    $assert_session->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForElementVisible('css', 'fieldset[data-drupal-selector="edit-source-configuration"]');
+    $this->assertNotEmpty($result);
     $assert_session->fieldExists('Test config value');
     $assert_session->fieldValueEquals('Test config value', 'This is default value.');
     $assert_session->fieldExists('Attribute 1');
diff --git a/core/modules/media/tests/src/FunctionalJavascript/MediaViewsWizardTest.php b/core/modules/media/tests/src/FunctionalJavascript/MediaViewsWizardTest.php
index 7b91c6e..58d3cd6 100644
--- a/core/modules/media/tests/src/FunctionalJavascript/MediaViewsWizardTest.php
+++ b/core/modules/media/tests/src/FunctionalJavascript/MediaViewsWizardTest.php
@@ -27,7 +27,8 @@ public function testMediaWizard() {
     $page->fillField('label', $view_id);
     $this->waitUntilVisible('.machine-name-value');
     $page->selectFieldOption('show[wizard_key]', 'media');
-    $assert_session->assertWaitOnAjaxRequest();
+    $result = $assert_session->waitForElementVisible('css', 'select[data-drupal-selector="edit-show-type"]');
+    $this->assertNotEmpty($result);
     $page->checkField('page[create]');
     $page->fillField('page[path]', $this->randomMachineName(16));
     $page->pressButton('Save and edit');
@@ -60,7 +61,10 @@ public function testMediaRevisionWizard() {
     $page->fillField('label', $view_id);
     $this->waitUntilVisible('.machine-name-value');
     $page->selectFieldOption('show[wizard_key]', 'media_revision');
-    $assert_session->assertWaitOnAjaxRequest();
+    // @TODO the ajax call at this point doesn't add or remove any element.
+    // Find a way to wait for this ajax to finish without using
+    // ::assertWaitOnAjaxRequest().
+    $session->wait(5000);
     $page->checkField('page[create]');
     $page->fillField('page[path]', $this->randomMachineName(16));
     $page->pressButton('Save and edit');