diff --git a/tests/src/FunctionalJavascript/EntityBrowserViewsWidgetTest.php b/tests/src/FunctionalJavascript/EntityBrowserViewsWidgetTest.php
index 5f4c8c1..e9c0df3 100644
--- a/tests/src/FunctionalJavascript/EntityBrowserViewsWidgetTest.php
+++ b/tests/src/FunctionalJavascript/EntityBrowserViewsWidgetTest.php
@@ -10,7 +10,7 @@ use Drupal\file\Entity\File;
  * @group entity_browser
  * @see \Drupal\entity_browser\Plugin\EntityBrowser\Widget\View
  */
-class EntityBrowserViewsWidgetTest extends EntityBrowserJavascriptTestBase {
+class EntityBrowserViewsWidgetTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * Modules to enable.
@@ -55,15 +55,15 @@ class EntityBrowserViewsWidgetTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->pageTextContains('example.jpg');
     $this->assertSession()->fieldExists($field);
     $this->getSession()->getPage()->fillField('filename', 'llama');
-    $this->getSession()->getPage()->pressButton('Apply');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Apply')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->fieldNotExists($field);
     $this->assertSession()->pageTextNotContains('example.jpg');
     $this->getSession()->getPage()->fillField('filename', 'example');
-    $this->getSession()->getPage()->pressButton('Apply');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Apply')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->pageTextContains('example.jpg');
-    $this->assertSession()->fieldExists($field);
+    $this->assertSession()->fieldExists($field)->check();
 
     // Test selection.
     $this->submitForm([
@@ -96,7 +96,7 @@ class EntityBrowserViewsWidgetTest extends EntityBrowserJavascriptTestBase {
     $this->drupalGet('/entity-browser/iframe/test_entity_browser_file', ['query' => ['uuid' => $uuid]]);
     $this->getSession()->getPage()->fillField('entity_browser_select[file:1]', TRUE);
     $this->getSession()->getPage()->fillField('entity_browser_select[file:2]', TRUE);
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     $this->assertSession()->pageTextContains('You can not select more than 1 entity.');
     $this->assertSession()->checkboxNotChecked('entity_browser_select[file:1]');
diff --git a/tests/src/FunctionalJavascript/EntityBrowserJavascriptTestBase.php b/tests/src/FunctionalJavascript/EntityBrowserWebDriverTestBase.php
similarity index 95%
rename from tests/src/FunctionalJavascript/EntityBrowserJavascriptTestBase.php
rename to tests/src/FunctionalJavascript/EntityBrowserWebDriverTestBase.php
index e54d99b..ff639b1 100644
--- a/tests/src/FunctionalJavascript/EntityBrowserJavascriptTestBase.php
+++ b/tests/src/FunctionalJavascript/EntityBrowserWebDriverTestBase.php
@@ -7,14 +7,20 @@ use Drupal\Core\Field\FieldStorageDefinitionInterface;
 use Drupal\file\Entity\File;
 use Drupal\field\Entity\FieldConfig;
 use Drupal\field\Entity\FieldStorageConfig;
-use Drupal\FunctionalJavascriptTests\JavascriptTestBase;
+use Drupal\FunctionalJavascriptTests\WebDriverTestBase;
+use Drupal\Tests\TestFileCreationTrait;
 
 /**
  * Base class for Entity browser Javascript functional tests.
  *
  * @package Drupal\Tests\entity_browser\FunctionalJavascript
  */
-abstract class EntityBrowserJavascriptTestBase extends JavascriptTestBase {
+abstract class EntityBrowserWebDriverTestBase extends WebDriverTestBase {
+
+  use TestFileCreationTrait {
+    getTestFiles as drupalGetTestFiles;
+    compareFiles as drupalCompareFiles;
+  }
 
   /**
    * Modules to enable.
@@ -210,7 +216,7 @@ abstract class EntityBrowserJavascriptTestBase extends JavascriptTestBase {
     $this->getSession()->getPage()->find('xpath', $xpathSelector)->click();
 
     if ($waitAfterAction) {
-      $this->waitForAjaxToFinish();
+      $this->assertSession()->assertWaitOnAjaxRequest();
     }
   }
 
diff --git a/tests/src/FunctionalJavascript/EntityReferenceWidgetTest.php b/tests/src/FunctionalJavascript/EntityReferenceWidgetTest.php
index 33a3086..ddbef21 100644
--- a/tests/src/FunctionalJavascript/EntityReferenceWidgetTest.php
+++ b/tests/src/FunctionalJavascript/EntityReferenceWidgetTest.php
@@ -14,7 +14,7 @@ use Drupal\user\Entity\Role;
  *
  * @group entity_browser
  */
-class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
+class EntityReferenceWidgetTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * {@inheritdoc}
@@ -90,12 +90,12 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     $this->drupalGet('/node/add/article');
     $page->fillField('title[0][value]', 'Referencing node 1');
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node1');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
-    $page->pressButton('Save');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->buttonExists('Save')->press();
 
     $assert_session->pageTextContains('Article Referencing node 1 has been created.');
     $nid = $this->container->get('entity.query')->get('node')->condition('title', 'Referencing node 1')->execute();
@@ -151,8 +151,8 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     $assert_session->buttonNotExists('edit-field-entity-reference1-current-items-0-replace-button');
 
     // Test the "Remove" button on the widget works.
-    $page->pressButton('Remove');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Remove')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->pageTextNotContains('Target example node 1');
 
     // Test the "Replace" button functionality.
@@ -199,30 +199,30 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     // Clicking on the button should empty the selection and automatically
     // open the browser again.
     $replace_button->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node3');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->wait(1000);
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     // Even in the AJAX-built markup for the newly selected element, the replace
     // button should be there.
     $assert_session->elementExists('css', 'input[data-drupal-selector="edit-field-entity-reference1-current-items-0-replace-button"]');
     // Adding a new node to the selection, however, should make it disappear.
     $open_iframe_link = $assert_session->elementExists('css', 'a[data-drupal-selector="edit-field-entity-reference1-entity-browser-entity-browser-link"]');
     $open_iframe_link->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node1');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->wait(1000);
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->elementNotExists('css', 'input[data-drupal-selector="edit-field-entity-reference1-current-items-0-replace-button"]');
-    $page->pressButton('Save');
+    $this->assertSession()->buttonExists('Save')->press();
     $assert_session->pageTextContains('Article Referencing node 1 has been updated.');
 
     // Test the replace button again with different field cardinalities.
@@ -231,14 +231,14 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     $page->fillField('title[0][value]', 'Referencing node 2');
     $open_iframe_link = $assert_session->elementExists('css', 'a[data-drupal-selector="edit-field-entity-reference1-entity-browser-entity-browser-link"]');
     $open_iframe_link->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node1');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->wait(1000);
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->elementContains('css', '#edit-field-entity-reference1-wrapper', 'Target example node 1');
     // All three buttons should be visible.
     $assert_session->elementExists('css', 'input[data-drupal-selector="edit-field-entity-reference1-current-items-0-remove-button"]');
@@ -247,14 +247,14 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     // Clicking on the button should empty the selection and automatically
     // open the browser again.
     $replace_button->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node2');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->wait(1000);
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->elementContains('css', '#edit-field-entity-reference1-wrapper', 'Referencing node 1');
 
     // Do the same as above but now with cardinality 2.
@@ -263,14 +263,14 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     $page->fillField('title[0][value]', 'Referencing node 3');
     $open_iframe_link = $assert_session->elementExists('css', 'a[data-drupal-selector="edit-field-entity-reference1-entity-browser-entity-browser-link"]');
     $open_iframe_link->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node1');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->wait(1000);
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->elementContains('css', '#edit-field-entity-reference1-wrapper', 'Target example node 1');
     // All three buttons should be visible.
     $assert_session->elementExists('css', 'input[data-drupal-selector="edit-field-entity-reference1-current-items-0-remove-button"]');
@@ -279,14 +279,14 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     // Clicking on the button should empty the selection and automatically
     // open the browser again.
     $replace_button->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node2');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->wait(1000);
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->elementContains('css', '#edit-field-entity-reference1-wrapper', 'Referencing node 1');
 
     // Verify that if the user cannot edit the entity, the "Edit" button does
@@ -297,15 +297,14 @@ class EntityReferenceWidgetTest extends EntityBrowserJavascriptTestBase {
     $this->drupalGet('node/add/article');
     $open_iframe_link = $assert_session->elementExists('css', 'a[data-drupal-selector="edit-field-entity-reference1-entity-browser-entity-browser-link"]');
     $open_iframe_link->click();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $session->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_node_view');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $page->checkField('edit-entity-browser-select-node1');
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $session->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $assert_session->buttonNotExists('edit-field-entity-reference1-current-items-0-edit-button');
-
   }
 
 }
diff --git a/tests/src/FunctionalJavascript/ImageFieldTest.php b/tests/src/FunctionalJavascript/ImageFieldTest.php
index 73e9877..1263d25 100644
--- a/tests/src/FunctionalJavascript/ImageFieldTest.php
+++ b/tests/src/FunctionalJavascript/ImageFieldTest.php
@@ -14,7 +14,7 @@ use Drupal\node\Entity\Node;
  *
  * @group entity_browser
  */
-class ImageFieldTest extends EntityBrowserJavascriptTestBase {
+class ImageFieldTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * Created file entity.
@@ -118,13 +118,13 @@ class ImageFieldTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->linkExists('Select images');
     $this->getSession()->getPage()->clickLink('Select images');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_view');
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $this->image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->getPage()->pressButton('Use selected');
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $this->image->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->assertSession()->pageTextContains('example.jpg');
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     // Check if the image thumbnail exists.
     $this->assertSession()->elementExists('xpath', '//*[@data-drupal-selector="edit-field-image-current-' . $this->image->id() . '-display"]');
     // Test if the image filename is present.
@@ -135,7 +135,7 @@ class ImageFieldTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->getPage()->fillField('field_image[current][1][meta][alt]', $alt_text);
     $this->getSession()->getPage()->fillField('field_image[current][1][meta][title]', $title_text);
     $this->getSession()->getPage()->fillField('title[0][value]', 'Node 1');
-    $this->getSession()->getPage()->pressButton('Save');
+    $this->assertSession()->buttonExists('Save')->press();
     $this->assertSession()->pageTextContains('Article Node 1 has been created.');
     $node = Node::load(1);
     $saved_alt = $node->get('field_image')[0]->alt;
@@ -144,9 +144,8 @@ class ImageFieldTest extends EntityBrowserJavascriptTestBase {
     $this->assertEquals($saved_title, $title_text);
     // Test the Delete functionality.
     $this->drupalGet('node/1/edit');
-    $this->assertSession()->buttonExists('Remove');
-    $this->getSession()->getPage()->pressButton('Remove');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Remove')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     // Image filename should not be present.
     $this->assertSession()->pageTextNotContains('example.jpg');
     $this->assertSession()->linkExists('Select entities');
@@ -157,16 +156,15 @@ class ImageFieldTest extends EntityBrowserJavascriptTestBase {
     $image2->save();
     \Drupal::service('file.usage')->add($image2, 'entity_browser', 'test', '1');
     $this->drupalGet('node/1/edit');
-    $this->assertSession()->buttonExists('Replace');
-    $this->getSession()->getPage()->pressButton('Replace');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Replace')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_iframe_view');
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image2->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->getPage()->pressButton('Use selected');
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image2->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->getSession()->wait(1000);
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     // Initial image should not be present, the new one should be there instead.
     $this->assertSession()->pageTextNotContains('example.jpg');
     $this->assertSession()->pageTextContains('example2.jpg');
@@ -190,7 +188,7 @@ class ImageFieldTest extends EntityBrowserJavascriptTestBase {
     // to allow png but the field widget is configured to allow jpg, so we
     // expect the field to override the widget.
     $this->getSession()->getPage()->attachFileToField('files[upload][]', $file_wrong_type);
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     if (version_compare(\Drupal::VERSION, '8.7', '>=')) {
       $this->assertSession()->responseContains('Only files with the following extensions are allowed: <em class="placeholder">jpg</em>.');
       $this->assertSession()->responseContains('The selected file <em class="placeholder">druplicon.png</em> cannot be uploaded.');
@@ -201,17 +199,17 @@ class ImageFieldTest extends EntityBrowserJavascriptTestBase {
     }
     // Upload an image bigger than the field widget's configured max size.
     $this->getSession()->getPage()->attachFileToField('files[upload][]', $file_too_big);
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->pageTextContains('The image was resized to fit within the maximum allowed dimensions of 40x40 pixels.');
     // Upload an image that passes validation and finish the upload.
     $this->getSession()->getPage()->attachFileToField('files[upload][]', $file_just_right);
-    $this->waitForAjaxToFinish();
-    $this->getSession()->getPage()->pressButton('Select files');
-    $this->getSession()->getPage()->pressButton('Use selected');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->buttonExists('Select files')->press();
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->assertSession()->pageTextContains('image-test.jpg');
     // Check that the file has uploaded to the correct sub-directory.
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $entity_id = $this->getSession()->evaluateScript('jQuery("#edit-field-image-wrapper [data-entity-id]").data("entity-id")');
     $this->assertStringStartsWith('file:', $entity_id);
     /** @var \Drupal\file\Entity\File $file */
diff --git a/tests/src/FunctionalJavascript/InlineEntityFormTest.php b/tests/src/FunctionalJavascript/InlineEntityFormTest.php
index 9f16fc0..6a3180b 100644
--- a/tests/src/FunctionalJavascript/InlineEntityFormTest.php
+++ b/tests/src/FunctionalJavascript/InlineEntityFormTest.php
@@ -11,7 +11,7 @@ use Behat\Mink\Element\NodeElement;
  *
  * @package Drupal\Tests\entity_browser\FunctionalJavascript
  */
-class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
+class InlineEntityFormTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * {@inheritdoc}
@@ -89,7 +89,7 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
     $page = $this->getSession()->getPage();
 
     $page->fillField('Title', 'Test IEF Title');
-    $page->pressButton('Add new Test File Media');
+    $this->assertSession()->buttonExists('Add new Test File Media')->press();
 
     $this->assertSession()->assertWaitOnAjaxRequest();
 
@@ -101,22 +101,22 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
     $page->checkField('entity_browser_select[file:1]');
     $page->checkField('entity_browser_select[file:2]');
 
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Use selected');
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->getSession()->switchToIFrame();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Create Test File Media');
+    $this->assertSession()->buttonExists('Create Test File Media')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Save');
+    $this->assertSession()->buttonExists('Save')->press();
 
     $this->drupalGet('node/1/edit');
     $page = $this->getSession()->getPage();
 
-    $page->pressButton('Edit');
+    $this->assertSession()->buttonExists('Edit')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Test reorder of elements.
@@ -124,10 +124,10 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
     $this->dragDropElement($dragged, 150, 0);
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Update Test File Media');
+    $this->assertSession()->buttonExists('Update Test File Media')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Edit');
+    $this->assertSession()->buttonExists('Edit')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Check that element on second position is test_file1 (file:1).
@@ -141,10 +141,10 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
     $this->click('input[name*="ief_media_type_file_field_remove_1_1"]');
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Update Test File Media');
+    $this->assertSession()->buttonExists('Update Test File Media')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Edit');
+    $this->assertSession()->buttonExists('Edit')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Check that remote button does not exist for already removed entity.
@@ -158,17 +158,17 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
       ->switchToIFrame('entity_browser_iframe_ief_entity_browser_file');
     $page->checkField('entity_browser_select[file:3]');
 
-    $page->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Use selected');
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->getSession()->switchToIFrame();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Update Test File Media');
+    $this->assertSession()->buttonExists('Update Test File Media')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Edit');
+    $this->assertSession()->buttonExists('Edit')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Check that element on second position is test_file3 (file:3).
@@ -188,14 +188,14 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
     $this->dragDropElement($dragged, 150, 0);
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Use selected');
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->getSession()->switchToIFrame();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Update Test File Media');
+    $this->assertSession()->buttonExists('Update Test File Media')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Edit');
+    $this->assertSession()->buttonExists('Edit')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Check that element on second position is test_file2 (file:2).
@@ -212,17 +212,17 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
       ->switchToIFrame('entity_browser_iframe_ief_entity_browser_file');
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('remove_3_0');
+    $this->assertSession()->buttonExists('remove_3_0')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Use selected');
+    $this->assertSession()->buttonExists('Use selected')->press();
     $this->getSession()->switchToIFrame();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Update Test File Media');
+    $this->assertSession()->buttonExists('Update Test File Media')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $page->pressButton('Edit');
+    $this->assertSession()->buttonExists('Edit')->press();
     $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Check that element on first position is test_file2 (file:2).
@@ -246,13 +246,12 @@ class InlineEntityFormTest extends EntityBrowserJavascriptTestBase {
     $page = $this->getSession()->getPage();
 
     $page->fillField('Title', 'Test IEF Title');
-    $page->pressButton('Add existing Test File Media');
+    $this->assertSession()->buttonExists('Add existing Test File Media')->press();
 
-    $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $this->getSession()
-      ->switchToIFrame('entity_browser_iframe_ief_entity_browser_file_modal');
+    $this->assertSession()->waitForElementVisible('css', 'iframe.entity_browser_iframe_image_browser_only_reference');
 
+    $this->createScreenshot(\Drupal::root() . '/sites/default/files/simpletest/screen.png');
     $this->assertSession()->pageTextContains('Test entity browser file modal');
   }
 
diff --git a/tests/src/FunctionalJavascript/MultiStepSelectionDisplayTest.php b/tests/src/FunctionalJavascript/MultiStepSelectionDisplayTest.php
index bfbd06a..9bfb5b6 100644
--- a/tests/src/FunctionalJavascript/MultiStepSelectionDisplayTest.php
+++ b/tests/src/FunctionalJavascript/MultiStepSelectionDisplayTest.php
@@ -11,7 +11,7 @@ use Drupal\Component\Utility\NestedArray;
  *
  * @package Drupal\Tests\entity_browser\FunctionalJavascript
  */
-class MultiStepSelectionDisplayTest extends EntityBrowserJavascriptTestBase {
+class MultiStepSelectionDisplayTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * Open iframe entity browser and change scope to iframe.
@@ -20,7 +20,7 @@ class MultiStepSelectionDisplayTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()
       ->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
   }
 
   /**
@@ -29,7 +29,7 @@ class MultiStepSelectionDisplayTest extends EntityBrowserJavascriptTestBase {
   protected function closeEntityBrowser() {
     $this->clickXpathSelector('//*[@data-drupal-selector="edit-use-selected"]');
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
   }
 
   /**
@@ -55,9 +55,9 @@ class MultiStepSelectionDisplayTest extends EntityBrowserJavascriptTestBase {
    * It's related to: Drupal.entityBrowserCommandQueue.executeCommands
    */
   protected function waitSelectionDisplayAjaxCommands() {
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->getSession()->wait(200);
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
   }
 
   /**
diff --git a/tests/src/FunctionalJavascript/PagerElementTest.php b/tests/src/FunctionalJavascript/PagerElementTest.php
index ad81f5d..97792f8 100644
--- a/tests/src/FunctionalJavascript/PagerElementTest.php
+++ b/tests/src/FunctionalJavascript/PagerElementTest.php
@@ -11,7 +11,7 @@ use Drupal\field\Entity\FieldStorageConfig;
  *
  * @group entity_browser
  */
-class PagerElementTest extends EntityBrowserJavascriptTestBase {
+class PagerElementTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * Modules to enable.
@@ -84,34 +84,35 @@ class PagerElementTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->buttonExists('‹ Previous')->hasAttribute('disabled');
     $this->assertSession()->pageTextContains('Page 1');
     $this->assertSession()->buttonExists('Next ›');
-    $this->getSession()->getPage()->pressButton('Next ›');
+    $this->assertSession()->buttonExists('Next ›')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 2.');
     $this->assertSession()->pageTextContains('Page 2');
-    $this->getSession()->getPage()->pressButton('Next ›');
+    $this->assertSession()->buttonExists('Next ›')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 3.');
     $this->assertSession()->pageTextContains('Page 3');
-    $this->getSession()->getPage()->pressButton('Next ›');
+    $this->assertSession()->buttonExists('Next ›')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 4.');
     $this->assertSession()->pageTextContains('Page 4');
     $this->assertSession()->buttonExists('Next ›')->hasAttribute('disabled');
 
     // Go back.
-    $this->getSession()->getPage()->pressButton('‹ Previous');
+    $this->assertSession()->buttonExists('‹ Previous')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 3.');
     $this->assertSession()->pageTextContains('Page 3');
-    $this->getSession()->getPage()->pressButton('‹ Previous');
+    $this->assertSession()->buttonExists('‹ Previous')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 2.');
     $this->assertSession()->pageTextContains('Page 2');
-    $this->getSession()->getPage()->pressButton('‹ Previous');
+    $this->assertSession()->buttonExists('‹ Previous')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 1.');
     $this->assertSession()->pageTextContains('Page 1');
     $this->assertSession()->buttonExists('‹ Previous')->hasAttribute('disabled');
 
     // Test reset button.
-    $this->getSession()->getPage()->pressButton('Last page');
+    $this->assertSession()->buttonExists('Last page')->press();
+
     $this->assertSession()->pageTextContains('Current page reported by the element is: 4.');
     $this->assertSession()->pageTextContains('Page 4');
-    $this->getSession()->getPage()->pressButton('First page');
+    $this->assertSession()->buttonExists('First page')->press();
     $this->assertSession()->pageTextContains('Current page reported by the element is: 1.');
     $this->assertSession()->pageTextContains('Page 1');
   }
diff --git a/tests/src/FunctionalJavascript/ParagraphsTest.php b/tests/src/FunctionalJavascript/ParagraphsTest.php
index b5ea4d8..b8dc63b 100644
--- a/tests/src/FunctionalJavascript/ParagraphsTest.php
+++ b/tests/src/FunctionalJavascript/ParagraphsTest.php
@@ -7,7 +7,7 @@ namespace Drupal\Tests\entity_browser\FunctionalJavascript;
  *
  * @group entity_browser
  */
-class ParagraphsTest extends EntityBrowserJavascriptTestBase {
+class ParagraphsTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * {@inheritdoc}
@@ -51,30 +51,30 @@ class ParagraphsTest extends EntityBrowserJavascriptTestBase {
     $title->setValue('Paragraph test');
 
     // Create a new paragraph referencing our article.
-    $this->getSession()->getPage()->pressButton('field_paragraph_content_embed_add_more');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('field_paragraph_content_embed_add_more')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_nodes');
-    $this->getSession()->getPage()->checkField('entity_browser_select[node:' . $node->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->fieldExists('entity_browser_select[node:' . $node->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Create another paragraph referencing our article.
-    $this->getSession()->getPage()->pressButton('field_paragraph_content_embed_add_more');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('field_paragraph_content_embed_add_more')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_nodes');
-    $this->getSession()->getPage()->checkField('entity_browser_select[node:' . $node->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->fieldExists('entity_browser_select[node:' . $node->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Ensure that two paragraphs now exist.
     $selector_1 = '.field--name-field-paragraph tbody tr:nth-child(1) .paragraphs-subform';
@@ -84,42 +84,43 @@ class ParagraphsTest extends EntityBrowserJavascriptTestBase {
 
     // Click the remove button on the second paragraph's article reference,
     // and ensure that only that instance of our article is removed.
-    $this->click('.field--name-field-paragraph tbody tr:nth-child(2) .paragraphs-subform [value="Remove"]');
-    $this->waitForAjaxToFinish();
+    $this->createScreenshot(\Drupal::root() . '/sites/default/files/simpletest/screen.png');
+    $this->assertSession()->elementExists('css', '.field--name-field-paragraph tbody tr:nth-child(2) .paragraphs-subform [value="Remove"]')->click();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->elementTextNotContains('css', $selector_2, 'Hello world');
     $this->assertSession()->elementTextContains('css', $selector_1, 'Hello world');
 
     // Testing nested paragraphs field.
-    $this->getSession()->getPage()->pressButton('field_paragraph_0_subform_field_paragraphs_nested_nested_paragraph_add_more');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('field_paragraph_0_subform_field_paragraphs_nested_nested_paragraph_add_more')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_nodes');
-    $this->getSession()->getPage()->checkField('entity_browser_select[node:' . $node_nested->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->fieldExists('entity_browser_select[node:' . $node_nested->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
-    $this->getSession()->getPage()->pressButton('field_paragraph_0_subform_field_paragraphs_nested_nested_paragraph_add_more');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('field_paragraph_0_subform_field_paragraphs_nested_nested_paragraph_add_more')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_nodes');
-    $this->getSession()->getPage()->checkField('entity_browser_select[node:' . $node_nested_second->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->fieldExists('entity_browser_select[node:' . $node_nested_second->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
     // Test editing node in nested paragraph.
     $this->assertSession()->elementTextContains('css', '.field--name-field-paragraphs-nested tbody tr:nth-child(1) .paragraphs-subform', 'Nested node');
     $this->click('.field--name-field-paragraphs-nested tbody tr:nth-child(1) .paragraphs-subform [value="Edit"]');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->elementExists('css', '.ui-dialog .node-article-edit-form .form-item-title-0-value [value="Nested node"]');
-    $this->getSession()->getPage()->pressButton('Close');
+    $this->assertSession()->buttonExists('Close')->press();
 
     // Submit the form.
     $this->submitForm([], 'Save');
@@ -134,7 +135,7 @@ class ParagraphsTest extends EntityBrowserJavascriptTestBase {
 
     // Test removing node in nested paragraph.
     $this->click('.field--name-field-paragraphs-nested tbody tr:nth-child(1) .paragraphs-subform [value="Remove"]');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->elementTextNotContains('css', '.field--name-field-paragraphs-nested tbody tr:nth-child(1) .paragraphs-subform', 'Nested node');
     $this->assertSession()->elementTextContains('css', '.field--name-field-paragraphs-nested tbody tr:nth-child(2) .paragraphs-subform', 'Nested second node');
 
diff --git a/tests/src/FunctionalJavascript/PluginsTest.php b/tests/src/FunctionalJavascript/PluginsTest.php
index 074e539..e9d787f 100644
--- a/tests/src/FunctionalJavascript/PluginsTest.php
+++ b/tests/src/FunctionalJavascript/PluginsTest.php
@@ -7,7 +7,7 @@ namespace Drupal\Tests\entity_browser\FunctionalJavascript;
  *
  * @group entity_browser
  */
-class PluginsTest extends EntityBrowserJavascriptTestBase {
+class PluginsTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * Tests the Entity browser iframe display plugin.
@@ -19,14 +19,12 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
 
     // Tests view widget on single display.
     $this->drupalGet('node/add/article');
-    $this->assertSession()->linkExists('Select entities');
-    $this->getSession()->getPage()->clickLink('Select entities');
-    $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
-
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->switchToIFrame();
-    $this->waitUntilVisible('.field--type-entity-reference .button');
+    $this->assertSession()->linkExists('Select entities')->click();
+   // $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']')->check();
+    $this->assertSession()->buttonExists('Select entities')->press();
+    $this->assertSession()->waitForElementVisible('css', '.field--type-entity-reference .button');
     $this->assertSession()->pageTextContains('lama.jpg');
 
     // Tests upload widget on single display. Gets the upload widget and sets
@@ -40,9 +38,8 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
 
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->switchToIFrame();
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']')->check();
+    $this->assertSession()->buttonExists('Select entities')->press();
     $this->waitUntilVisible('.field--type-entity-reference .button');
     $this->assertSession()->pageTextContains('lama.jpg');
 
@@ -57,9 +54,8 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->linkExists('upload');
 
     $this->clickLink('view');
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->switchToIFrame();
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']')->check();
+    $this->assertSession()->buttonExists('Select entities')->press();
     $this->waitUntilVisible('.field--type-entity-reference .button');
     $this->assertSession()->pageTextContains('lama.jpg');
 
@@ -69,11 +65,10 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
     $this->clickLink('upload');
     $this->getSession()->getPage()->attachFileToField('edit-upload-upload', $this->container->get('file_system')->realpath($image->getFileUri()));
-    $this->waitForAjaxToFinish();
-    $this->getSession()->getPage()->checkField('upload[file_2][selected]');
-    $this->getSession()->getPage()->pressButton('Select files');
-    $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->fieldExists('upload[file_2][selected]');
+    $this->assertSession()->buttonExists('Select files')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->pageTextContains('lama.jpg');
     // Tests view widget with drop down widget selector.
     $this->getEntityBrowser('test_entity_browser_file', 'iframe', 'drop_down', 'no_display');
@@ -91,9 +86,8 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
     $this->assertSession()->selectExists('edit-widget');
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->switchToIFrame();
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']')->check();
+    $this->assertSession()->buttonExists('Select entities')->press();
     $this->waitUntilVisible('.field--type-entity-reference .button');
     $this->assertSession()->pageTextContains('lama.jpg');
 
@@ -103,13 +97,12 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
     $this->getSession()->getPage()->selectFieldOption('edit-widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->getSession()->getPage()->attachFileToField('files[upload][]', $this->container->get('file_system')->realpath($image->getFileUri()));
-    $this->waitForAjaxToFinish();
-    $this->getSession()->getPage()->checkField('upload[file_3][selected]');
-    $this->getSession()->getPage()->pressButton('Select files');
-    $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->fieldExists('upload[file_3][selected]');
+    $this->assertSession()->buttonExists('Select files')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     // In iframe I get page not found, so this fails.
     $this->assertSession()->pageTextContains('lama.jpg');
     // Tests view selection display.
@@ -149,27 +142,21 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
 
     $this->getSession()->getPage()->attachFileToField('files[upload][]', $this->container->get('file_system')->realpath($image1->getFileUri()));
-    $this->waitForAjaxToFinish();
-    $this->getSession()->getPage()->pressButton('Select files');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->buttonExists('Select files')->press();
 
     $this->assertSession()->pageTextContains('first_file.jpg');
     $this->assertSession()->pageTextNotContains('second_file.jpg');
 
     $this->getSession()->getPage()->clickLink('view');
-    // This shows page not found,
-    // which is caused by https://www.drupal.org/node/2771547
+    // This shows page not found, which is caused by https://www.drupal.org/node/2771547
     // Uncomment this hunk when a fix for that problem lands.
-    // $this->getSession()
-    // ->getPage()
-    // ->checkField('entity_browser_select[file:' . $image2->id() . ']');
-    // $this->getSession()->getPage()->pressButton('Select entities');
-    // $this->assertSession()
-    // ->responseContains('edit-selected-items-2-1-remove-button');
-    // $this->assertSession()
-    // ->responseContains('edit-selected-items-1-0-remove-button');
-    // $this->getSession()->getPage()->pressButton('Use selected');
-    // $this->getSession()->switchToIFrame();
-    // $this->waitForAjaxToFinish();
+    // $this->assertSession()->fieldExists('entity_browser_select[file:' . $image2->id() . ']')->check();
+    // $this->assertSession()->buttonExists('Select entities')->press();
+    // $this->assertSession()->responseContains('edit-selected-items-2-1-remove-button');
+    // $this->assertSession()->responseContains('edit-selected-items-1-0-remove-button');
+    // $this->assertSession()->buttonExists('Use selected')->press();
+    // $this->assertSession()->assertWaitOnAjaxRequest();
     // $this->assertSession()->pageTextContains('first_file.jpg');
     // $this->assertSession()->pageTextContains('second_file.jpg');
   }
@@ -188,14 +175,12 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $image = $this->createFile('lama');
 
     $this->drupalGet('node/add/article');
-    $this->assertSession()->buttonExists('Select entities');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
-
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
-    $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Select entities')->press();
+    //$this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']')->check();
+    $this->assertSession()->buttonExists('Select entities')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
     $this->assertSession()->pageTextContains('lama.jpg');
   }
@@ -214,10 +199,10 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->drupalGet('test');
 
     $this->assertSession()->buttonExists('Select entities');
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
-    $this->getSession()->getPage()->checkField('entity_browser_select[file:' . $image->id() . ']');
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']')->check();
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // TODO test if entities were selected. Will most likely need a custom event
     // subscriber that displays a message or something along those lines.
diff --git a/tests/src/FunctionalJavascript/UploadWidgetTest.php b/tests/src/FunctionalJavascript/UploadWidgetTest.php
index b7bf839..4610583 100644
--- a/tests/src/FunctionalJavascript/UploadWidgetTest.php
+++ b/tests/src/FunctionalJavascript/UploadWidgetTest.php
@@ -9,7 +9,7 @@ use Drupal\user\Entity\Role;
  *
  * @group entity_browser
  */
-class UploadWidgetTest extends EntityBrowserJavascriptTestBase {
+class UploadWidgetTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * {@inheritdoc}
@@ -42,10 +42,10 @@ class UploadWidgetTest extends EntityBrowserJavascriptTestBase {
     // Go to the widget standalone page and test the upload.
     $this->drupalGet($browser->getDisplay()->path());
     $page->attachFileToField('edit-upload-upload', \Drupal::root() . '/core/misc/druplicon.png');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->fieldExists('druplicon.png');
-    $page->pressButton('Select files');
-    $this->assertSession()->statusCodeEquals(200);
+    $this->assertSession()->buttonExists('Select files')->press();
+    //$this->assertSession()->statusCodeEquals(200);
 
     // Check if the file was correctly uploaded to the EB destination.
     $this->assertFileExists('public://druplicon.png');
@@ -60,10 +60,10 @@ class UploadWidgetTest extends EntityBrowserJavascriptTestBase {
 
     $this->drupalGet($browser->getDisplay()->path());
     $page->attachFileToField('edit-upload-upload', \Drupal::root() . '/core/misc/druplicon.png');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->fieldExists('druplicon.png');
-    $page->pressButton('Fancy submit');
-    $this->assertSession()->statusCodeEquals(200);
+    $this->assertSession()->buttonExists('Fancy submit')->press();
+    //$this->assertSession()->statusCodeEquals(200);
 
     // Check if the file was correctly uploaded to the EB destination.
     $this->assertFileExists('public://some_location/druplicon.png');