diff --git a/modules/entity_form/src/Tests/InlineEntityIntegrationTest.php b/modules/entity_form/src/Tests/InlineEntityIntegrationTest.php
deleted file mode 100644
index 6ae5b79..0000000
--- a/modules/entity_form/src/Tests/InlineEntityIntegrationTest.php
+++ /dev/null
@@ -1,76 +0,0 @@
-<?php
-
-namespace Drupal\entity_browser_entity_form\Tests;
-
-use Drupal\simpletest\WebTestBase;
-
-/**
- * Tests integration with Inline entity form.
- *
- * @group entity_browser_entity_form
- */
-class InlineEntityIntegrationTest extends WebTestBase {
-
-  /**
-   * Modules to enable.
-   *
-   * @var array
-   */
-  public static $modules = [
-    'entity_browser_entity_form',
-    'node',
-    'field_ui',
-    'entity_browser_entity_form_test',
-  ];
-
-  /**
-   * {@inheritdoc}
-   */
-  protected function setUp() {
-    parent::setUp();
-    // Delete unnecessary entity browser.
-    $browser = $this->container->get('entity_type.manager')->getStorage('entity_browser')->load('entity_browser_test_entity_form');
-    $this->container->get('entity_type.manager')->getStorage('entity_browser')->delete([$browser]);
-  }
-
-  /**
-   * Tests integration with Inline entity form.
-   */
-  public function testInlineEntityIntegration() {
-    $account = $this->drupalCreateUser([
-      'administer node form display',
-      'administer node display',
-      'create article content',
-    ]);
-    $this->drupalLogin($account);
-    $this->drupalGet('admin/structure/types/manage/article/form-display');
-    $edit = [
-      'fields[field_content_reference][region]' => 'content',
-      'fields[field_content_reference][type]' => 'inline_entity_form_complex',
-    ];
-    $this->drupalPostForm(NULL, $edit, t('Save'));
-    $this->drupalPostAjaxForm(NULL, [], 'field_content_reference_settings_edit');
-    $this->assertRaw('fields[field_content_reference][settings_edit_form][third_party_settings][entity_browser_entity_form][entity_browser_id]', 'Field to select entity browser is available.');
-    $edit = [
-      'fields[field_content_reference][settings_edit_form][third_party_settings][entity_browser_entity_form][entity_browser_id]' => 'entity_browser_entity_form_test',
-      'fields[field_content_reference][settings_edit_form][settings][allow_existing]' => TRUE,
-    ];
-    $this->drupalPostAjaxForm(NULL, $edit, 'field_content_reference_plugin_settings_update');
-    $this->drupalPostForm(NULL, [], t('Save'));
-    $this->assertText('Entity browser: Entity browser entity form test', 'Settings summary is working correctly.');
-
-    $this->drupalGet('node/add');
-    $elements = $this->xpath('//input[@type="submit" and @value="Add existing node"]');
-    $button_name = $elements[0]->attributes()['name'];
-    $this->drupalPostAjaxForm(NULL, [], $button_name);
-    $this->assertLink('Select entities', 0, 'Entity browser is available.');
-
-    $browsers = $this->container->get('entity_type.manager')->getStorage('entity_browser')->loadMultiple();
-    $browser = current($browsers);
-    $this->container->get('entity_type.manager')->getStorage('entity_browser')->delete([$browser]);
-    $this->drupalGet('admin/structure/types/manage/article/form-display');
-    $this->drupalPostAjaxForm(NULL, [], 'field_content_reference_settings_edit');
-    $this->assertText(t('There are no entity browsers available. You can create one here'), 'Massage displays when no entity browser is available.');
-  }
-
-}
diff --git a/modules/entity_form/src/tests/FunctionalJavascript/InlineEntityIntegrationTest.php b/modules/entity_form/src/tests/FunctionalJavascript/InlineEntityIntegrationTest.php
new file mode 100644
index 0000000..1141dc1
--- /dev/null
+++ b/modules/entity_form/src/tests/FunctionalJavascript/InlineEntityIntegrationTest.php
@@ -0,0 +1,90 @@
+<?php
+
+namespace Drupal\entity_browser_entity_form\FunctionalJavascript\Tests;
+
+use Drupal\FunctionalJavascriptTests\WebDriverTestBase;
+
+/**
+ * Tests integration with Inline entity form.
+ *
+ * @group entity_browser_entity_form
+ */
+class InlineEntityIntegrationTest extends WebDriverTestBase {
+
+  /**
+   * The entity browser storage.
+   *
+   * @var \Drupal\Core\Entity\EntityStorageInterface
+   */
+  protected $browserStorage;
+
+  /**
+   * Modules to enable.
+   *
+   * @var array
+   */
+  public static $modules = [
+    'entity_browser_entity_form',
+    'node',
+    'field_ui',
+    'entity_browser_entity_form_test',
+  ];
+
+  /**
+   * {@inheritdoc}
+   */
+  protected function setUp() {
+    parent::setUp();
+    $this->entityBrowserStorage = $this->container
+      ->get('entity_type.manager')
+      ->getStorage('entity_browser');
+    $this->browserStorage = $this->container
+      ->get('entity_type.manager')
+      ->getStorage('entity_browser');
+    // Delete unnecessary entity browser.
+    $entityBrowserToDelete = $this->browserStorage
+      ->load('entity_browser_test_entity_form');
+    $this->browserStorage->delete([$entityBrowserToDelete]);
+  }
+
+  /**
+   * Tests integration with Inline entity form.
+   */
+  public function testInlineEntityIntegration() {
+    $account = $this->drupalCreateUser([
+      'administer node form display',
+      'administer node display',
+      'create article content',
+    ]);
+    $this->drupalLogin($account);
+    $this->drupalGet('admin/structure/types/manage/article/form-display');
+    $edit = [
+      'fields[field_content_reference][region]' => 'content',
+      'fields[field_content_reference][type]' => 'inline_entity_form_complex',
+    ];
+    $this->submitForm($edit, 'Save');
+    $this->submitForm([], 'field_content_reference_settings_edit');
+    $this->assertSession()->responseContains('fields[field_content_reference][settings_edit_form][third_party_settings][entity_browser_entity_form][entity_browser_id]', 'Field to select entity browser is available.');
+    $edit = [
+      'fields[field_content_reference][settings_edit_form][third_party_settings][entity_browser_entity_form][entity_browser_id]' => 'entity_browser_entity_form_test',
+      'fields[field_content_reference][settings_edit_form][settings][allow_existing]' => TRUE,
+    ];
+    $this->submitForm($edit, 'field_content_reference_plugin_settings_update');
+    $this->submitForm([], 'Save');
+    $this->assertSession()->responseContains('Entity browser: Entity browser entity form test', 'Settings summary is working correctly.');
+
+    $this->drupalGet('node/add');
+    $this->assertSession()->buttonExists('Add existing node')->press();
+    $this->assertSession()->responseContains('Select entities', 'Entity browser is available.');
+
+    $browsers = $this->browserStorage->loadMultiple();
+    $browser = current($browsers);
+    $this->browserStorage->delete([$browser]);
+    $this->drupalGet('admin/structure/types/manage/article/form-display');
+    $this->submitForm([], 'field_content_reference_settings_edit');
+    $this->assertSession()->selectExists('fields[field_content_reference][type]')->setValue('entity_browser_entity_reference');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->responseContains('There are no entity browsers available. You can create one here', 'Massage displays when no entity browser is available.');
+  }
+
+}
diff --git a/modules/entity_form/tests/src/FunctionalJavascript/EntityFormWidgetTest.php b/modules/entity_form/tests/src/FunctionalJavascript/EntityFormWidgetTest.php
index d1c603b..75bfcfd 100644
--- a/modules/entity_form/tests/src/FunctionalJavascript/EntityFormWidgetTest.php
+++ b/modules/entity_form/tests/src/FunctionalJavascript/EntityFormWidgetTest.php
@@ -5,14 +5,14 @@ namespace Drupal\Tests\entity_browser_entity_form\FunctionalJavascript;
 use Drupal\Core\Field\FieldStorageDefinitionInterface;
 use Drupal\field\Entity\FieldConfig;
 use Drupal\field\Entity\FieldStorageConfig;
-use Drupal\FunctionalJavascriptTests\JavascriptTestBase;
+use Drupal\FunctionalJavascriptTests\WebDriverTestBase;
 
 /**
  * Class for Entity browser entity form Javascript functional tests.
  *
  * @group entity_browser_entity_form
  */
-class EntityFormWidgetTest extends JavascriptTestBase {
+class EntityFormWidgetTest extends WebDriverTestBase {
 
   /**
    * Modules to enable.
@@ -102,12 +102,12 @@ class EntityFormWidgetTest extends JavascriptTestBase {
     $this->assertSession()->buttonNotExists('Save entity');
     $this->assertSession()->buttonExists('Save node');
 
-    // Make sure that the widget works correctly with the field widget
+    // Make sure that the widget works correctly with the field widget.
     $this->drupalGet('node/add/foo');
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_entity_browser_test_entity_form');
     $this->getSession()->getPage()->fillField('inline_entity_form[title][0][value]', 'War is peace');
-    $this->getSession()->getPage()->pressButton('Save node');
+    $this->assertSession()->buttonExists('Save node')->press();
 
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
@@ -115,7 +115,7 @@ class EntityFormWidgetTest extends JavascriptTestBase {
 
     $this->assertSession()->pageTextContains('War is peace');
     $this->getSession()->getPage()->fillField('title[0][value]', 'Freedom is slavery');
-    $this->getSession()->getPage()->pressButton('Save');
+    $this->assertSession()->buttonExists('Save')->press();
 
     $parent_node = $this->container->get('entity_type.manager')
       ->getStorage('node')
@@ -131,8 +131,8 @@ class EntityFormWidgetTest extends JavascriptTestBase {
     $this->getSession()->getPage()->clickLink('Select entities');
     $this->getSession()->switchToIFrame('entity_browser_iframe_entity_browser_test_entity_form');
     $this->getSession()->getPage()->fillField('inline_entity_form[title][0][value]', 'War is peace');
-    $this->getSession()->getPage()->pressButton('Save node');
-    $this->getSession()->getPage()->pressButton('Use selected');
+    $this->assertSession()->buttonExists('Save node')->press();
+    $this->assertSession()->buttonExists('Use selected')->press();
 
     // Switch back to the main page.
     $this->getSession()->switchToIFrame();
@@ -140,7 +140,7 @@ class EntityFormWidgetTest extends JavascriptTestBase {
 
     $this->assertSession()->pageTextContains('War is peace');
     $this->getSession()->getPage()->fillField('title[0][value]', 'Ignorance is strength');
-    $this->getSession()->getPage()->pressButton('Save');
+    $this->assertSession()->buttonExists('Save')->press();
 
     $parent_node = $this->container->get('entity_type.manager')
       ->getStorage('node')
diff --git a/src/Tests/EntityBrowserUITest.php b/src/Tests/EntityBrowserUITest.php
index e7e5e50..816d5b1 100644
--- a/src/Tests/EntityBrowserUITest.php
+++ b/src/Tests/EntityBrowserUITest.php
@@ -3,14 +3,19 @@
 namespace Drupal\entity_browser\Tests;
 
 use Drupal\file\Entity\File;
-use Drupal\simpletest\WebTestBase;
+use Drupal\Tests\image\Functional\ImageFieldTestBase;
+use Drupal\Tests\TestFileCreationTrait;
 
 /**
  * Tests the entity browser UI.
  *
  * @group entity_browser
  */
-class EntityBrowserUITest extends WebTestBase {
+class EntityBrowserUITest extends ImageFieldTestBase {
+
+  use TestFileCreationTrait {
+    getTestFiles as drupalGetTestFiles;
+  }
 
   /**
    * Modules to enable.
@@ -34,14 +39,14 @@ class EntityBrowserUITest extends WebTestBase {
     $this->drupalLogin($account);
     // Go to the entity browser iframe link.
     $this->drupalGet('/entity-browser/iframe/test_entity_browser_iframe');
-    $this->assertRaw('Select');
+    $this->assertSession()->responseContains('Select');
     $this->drupalGet('/admin/config/content/entity_browser/test_entity_browser_iframe/widgets');
     $edit = [
       'table[871dbf77-012e-41cb-b32a-ada353d2de35][form][submit_text]' => 'Different',
     ];
-    $this->drupalPostForm(NULL, $edit, 'Save');
+    $this->submitForm($edit, 'Finish');
     $this->drupalGet('/entity-browser/iframe/test_entity_browser_iframe');
-    $this->assertRaw('Different');
+    $this->assertSession()->responseContains('Different');
   }
 
   /**
@@ -55,16 +60,16 @@ class EntityBrowserUITest extends WebTestBase {
     $this->drupalLogin($account);
     // Go to the entity browser iframe link.
     $this->drupalGet('/entity-browser/iframe/test_entity_browser_token');
-    $image = current($this->drupalGetTestFiles('image'));
+    $image = current($this->getTestFiles('image'));
     $edit = [
       'files[upload][]' => $this->container->get('file_system')->realpath($image->uri),
     ];
-    $this->drupalPostForm(NULL, $edit, 'Select files');
+    $this->submitForm($edit, 'Select files');
 
     $file = File::load(1);
     // Test entity browser token that has upload location configured to
     // public://[current-user:account-name]/.
-    $this->assertEqual($file->getFileUri(), 'public://' . $account->getUsername() . '/' . $file->getFilename(), 'Image has the correct uri.');
+    $this->assertEquals($file->getFileUri(), 'public://' . $account->getUsername() . '/' . $file->getFilename(), 'Image has the correct uri.');
   }
 
 }
diff --git a/src/Tests/FormElementTest.php b/src/Tests/FormElementTest.php
index 0a4ac41..15fd30f 100644
--- a/src/Tests/FormElementTest.php
+++ b/src/Tests/FormElementTest.php
@@ -3,14 +3,14 @@
 namespace Drupal\entity_browser\Tests;
 
 use Drupal\entity_browser\Element\EntityBrowserElement;
-use Drupal\simpletest\WebTestBase;
+use Drupal\Tests\BrowserTestBase;
 
 /**
  * Tests the entity browser form element.
  *
  * @group entity_browser
  */
-class FormElementTest extends WebTestBase {
+class FormElementTest extends BrowserTestBase {
 
   /**
    * Modules to enable.
@@ -22,7 +22,7 @@ class FormElementTest extends WebTestBase {
   /**
    * Test nodes.
    *
-   * @var \Drupal\node\NodeInterface[]
+   * @var array
    */
   protected $nodes;
 
@@ -48,27 +48,32 @@ class FormElementTest extends WebTestBase {
    */
   public function testFormElement() {
     $this->drupalGet('/test-element');
-    $this->assertLink('Select entities', 0, 'Trigger link found.');
-    $this->assertFieldByXPath("//input[@type='hidden' and @id='edit-fancy-entity-browser-target']", '', "Entity browser's hidden element found.");
+    $this->assertSession()->linkExists('Select entities', 0, 'Trigger link found.');
 
-    $edit = [
-      'fancy_entity_browser[entity_ids]' => $this->nodes[0]->getEntityTypeId() . ':' . $this->nodes[0]->id() . ' ' . $this->nodes[1]->getEntityTypeId() . ':' . $this->nodes[1]->id(),
+    $ids = [
+      $this->nodes[0]->getEntityTypeId() . ':' . $this->nodes[0]->id(),
+      $this->nodes[1]->getEntityTypeId() . ':' . $this->nodes[1]->id()
     ];
-    $this->drupalPostForm(NULL, $edit, 'Submit');
+
+    $ids = implode(' ', $ids);
+
+    $this->assertSession()->hiddenFieldExists("fancy_entity_browser[entity_ids]")->setValue($ids);
+
+    $this->assertSession()->buttonExists('Submit')->press();
     $expected = 'Selected entities: ' . $this->nodes[0]->label() . ', ' . $this->nodes[1]->label();
-    $this->assertText($expected, 'Selected entities detected.');
+    $this->assertSession()->responseContains($expected, 'Selected entities detected.');
 
     $default_entity = $this->nodes[0]->getEntityTypeId() . ':' . $this->nodes[0]->id();
     $this->drupalGet('/test-element', ['query' => ['default_entity' => $default_entity, 'selection_mode' => EntityBrowserElement::SELECTION_MODE_EDIT]]);
-    $this->assertLink('Select entities', 0, 'Trigger link found.');
-    $this->assertFieldByXPath("//input[@type='hidden' and @id='edit-fancy-entity-browser-target']", $default_entity, "Entity browser's hidden element found.");
+    $this->assertSession()->linkExists('Select entities', 0, 'Trigger link found.');
+    $this->assertSession()->fieldValueEquals("fancy_entity_browser[target]", $default_entity);
 
     $edit = [
       'fancy_entity_browser[entity_ids]' => $this->nodes[1]->getEntityTypeId() . ':' . $this->nodes[1]->id() . ' ' . $this->nodes[0]->getEntityTypeId() . ':' . $this->nodes[0]->id(),
     ];
-    $this->drupalPostForm(NULL, $edit, 'Submit');
+    $this->submitForm($edit, 'Submit');
     $expected = 'Selected entities: ' . $this->nodes[1]->label() . ', ' . $this->nodes[0]->label();
-    $this->assertText($expected, 'Selected entities detected.');
+    $this->assertSession()->responseContains($expected, 'Selected entities detected.');
   }
 
 }
diff --git a/tests/src/FunctionalJavascript/EntityBrowserTest.php b/tests/src/FunctionalJavascript/EntityBrowserTest.php
index 857d2c1..b4c4389 100644
--- a/tests/src/FunctionalJavascript/EntityBrowserTest.php
+++ b/tests/src/FunctionalJavascript/EntityBrowserTest.php
@@ -7,7 +7,7 @@ namespace Drupal\Tests\entity_browser\FunctionalJavascript;
  *
  * @group entity_browser
  */
-class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
+class EntityBrowserTest extends EntityBrowserWebDriverTestBase {
 
   /**
    * Tests single widget selector.
@@ -30,22 +30,22 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
 
-    $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
+    //$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->assertSession()->fieldExists('entity_browser_select[file:' . $image->id() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // Switch back to the main page.
-    $this->getSession()->switchToIFrame();
-    $this->waitForAjaxToFinish();
+    //$this->getSession()->switchToIFrame();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     // Test the Edit functionality.
     $this->assertSession()->pageTextContains('llama.jpg');
     $this->assertSession()->buttonExists('Edit');
     // @TODO Test the edit button.
     // Test the Delete functionality.
     $this->assertSession()->buttonExists('Remove');
-    $this->getSession()->getPage()->pressButton('Remove');
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->buttonExists('Remove')->press();
+    $this->assertSession()->assertWaitOnAjaxRequest();
     $this->assertSession()->pageTextNotContains('llama.jpg');
     $this->assertSession()->linkExists('Select entities');
   }
@@ -63,7 +63,7 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
       ->setWidgetSelector('tabs');
     $browser->save();
 
-    $this->assertEquals($browser->getWidgetSelector()->getPluginId(), 'tabs', 'Widget selector is set to tabs.');
+    $this->assertEquals('tabs', $browser->getWidgetSelector()->getPluginId(), 'Widget selector is set to tabs.');
 
     // Create a file.
     $image = $this->createFile('llama');
@@ -76,7 +76,7 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
 
-    $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
+    //$this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
 
     $this->assertSession()->linkExists('dummy');
     $this->assertSession()->linkExists('view');
@@ -84,11 +84,11 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
 
     $this->assertEquals('is-active active', $this->getSession()->getPage()->findLink('view')->getAttribute('class'));
 
-    $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() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
+    //$this->getSession()->switchToIFrame();
 
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
     $this->assertSession()->pageTextContains('llama.jpg');
 
@@ -99,11 +99,10 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     // This is producing an error. Still investigating
     // InvalidStateError: DOM Exception 11: An attempt was made to use an object
     // that is not, or is no longer, usable.
-    //$edit = [
+    // $edit = [
     //  'files[upload][]' => $this->container->get('file_system')->realpath($image2->getFileUri()),
-    //];
-    // $this->drupalPostForm(NULL, $edit, 'Select files');.
-
+    // ];
+    // $this->submitForm($edit, 'Select files');.
     \Drupal::state()->set('eb_test_dummy_widget_access', FALSE);
     $this->drupalGet('entity-browser/iframe/test_entity_browser_file');
     $this->assertSession()->linkNotExists('dummy');
@@ -148,20 +147,23 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     $this->assertSession()->linkExists('Select entities');
     $this->getSession()->getPage()->clickLink('Select entities');
 
-    $this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
+    //$this->getSession()->switchToIFrame('entity_browser_iframe_test_entity_browser_file');
 
     $this->assertSession()->selectExists('widget');
-    $this->assertSession()->optionExists('widget', 'cbc59500-04ab-4395-b063-c561f0e3bf80'); // Dummy
-    $this->assertSession()->optionExists('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87'); // Upload
-    $this->assertSession()->optionExists('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d'); // View
+    // Dummy.
+    $this->assertSession()->optionExists('widget', 'cbc59500-04ab-4395-b063-c561f0e3bf80');
+    // Upload.
+    $this->assertSession()->optionExists('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87');
+    // View.
+    $this->assertSession()->optionExists('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d');
     // Selects the view widget.
     $this->getSession()->getPage()->selectFieldOption('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d');
 
-    $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() . ']');
+    $this->assertSession()->buttonExists('Select entities')->press();
+    //$this->getSession()->switchToIFrame();
 
-    $this->waitForAjaxToFinish();
+    $this->assertSession()->assertWaitOnAjaxRequest();
 
     $this->assertSession()->pageTextContains('llama.jpg');
 
@@ -172,12 +174,14 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     // Causes a fatal.
     // Selects the upload widget.
     // $this->getSession()->getPage()->selectFieldOption('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87');.
-
     \Drupal::state()->set('eb_test_dummy_widget_access', FALSE);
     $this->drupalGet('entity-browser/iframe/test_entity_browser_file');
-    $this->assertSession()->optionNotExists('widget', 'cbc59500-04ab-4395-b063-c561f0e3bf80'); // Dummy
-    $this->assertSession()->optionExists('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87'); // Upload
-    $this->assertSession()->optionExists('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d'); // View
+    // Dummy.
+    $this->assertSession()->optionNotExists('widget', 'cbc59500-04ab-4395-b063-c561f0e3bf80');
+    // Upload.
+    $this->assertSession()->optionExists('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87');
+    // View.
+    $this->assertSession()->optionExists('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d');
     $this->assertSession()->responseHeaderContains('X-Drupal-Cache-Contexts', 'eb_dummy');
 
     // Move dummy widget to the first place and make sure it does not appear.
@@ -187,9 +191,12 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     $browser->getWidget('cbc59500-04ab-4395-b063-c561f0e3bf80')->setWeight(-15);
     $browser->save();
     $this->drupalGet('entity-browser/iframe/test_entity_browser_file');
-    $this->assertSession()->optionNotExists('widget', 'cbc59500-04ab-4395-b063-c561f0e3bf80'); // Dummy
-    $this->assertSession()->optionExists('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87'); // Upload
-    $this->assertSession()->optionExists('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d'); // View
+    // Dummy.
+    $this->assertSession()->optionNotExists('widget', 'cbc59500-04ab-4395-b063-c561f0e3bf80');
+    // Upload.
+    $this->assertSession()->optionExists('widget', '2dc1ab07-2f8f-42c9-aab7-7eef7f8b7d87');
+    // View.
+    $this->assertSession()->optionExists('widget', '774798f1-5ec5-4b63-84bd-124cd51ec07d');
     $this->assertSession()->pageTextNotContains('This is dummy widget.');
     $this->assertSession()->responseHeaderContains('X-Drupal-Cache-Contexts', 'eb_dummy');
   }
@@ -242,13 +249,13 @@ class EntityBrowserTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->switchToIFrame('entity_browser_iframe_multiple_submit_example');
 
     // Click the second submit button to make sure the widget does not close.
-    $this->getSession()->getPage()->pressButton('Second submit button');
+    $this->assertSession()->buttonExists('Second submit button')->press();
 
     // Check that the entity browser widget is still open.
     $this->getSession()->getPage()->hasButton('Second submit button');
 
     // Click the primary submit button to close the widget.
-    $this->getSession()->getPage()->pressButton('Select entities');
+    $this->assertSession()->buttonExists('Select entities')->press();
 
     // Check that the entity browser widget is closed.
     $this->assertSession()->buttonNotExists('Second submit button');
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 36955f8..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,13 @@ 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 5e51f4a..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.
@@ -44,7 +44,6 @@ class PagerElementTest extends EntityBrowserJavascriptTestBase {
       ],
     ])->save();
 
-
     FieldConfig::create([
       'field_name' => 'field_reference_pager',
       'entity_type' => 'node',
@@ -85,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 d94946c..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,8 +142,8 @@ 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');
@@ -158,15 +151,14 @@ class PluginsTest extends EntityBrowserJavascriptTestBase {
     $this->getSession()->getPage()->clickLink('view');
     // 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()->pageTextContains('first_file.jpg');
-    //$this->assertSession()->pageTextContains('second_file.jpg');
+    // $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');
   }
 
   /**
@@ -183,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');
   }
@@ -209,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');
diff --git a/tests/src/Kernel/Extension/EntityBrowserTest.php b/tests/src/Kernel/Extension/EntityBrowserTest.php
index 9be128c..3a63410 100644
--- a/tests/src/Kernel/Extension/EntityBrowserTest.php
+++ b/tests/src/Kernel/Extension/EntityBrowserTest.php
@@ -13,7 +13,6 @@ use Drupal\entity_browser\WidgetInterface;
 use Drupal\entity_browser\WidgetSelectorInterface;
 use Drupal\entity_browser\SelectionDisplayInterface;
 use Drupal\KernelTests\KernelTestBase;
-use Drupal\user\Entity\User;
 use Drupal\views\Entity\View;
 
 /**