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 7591391..6cfe2aa 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.
@@ -103,12 +103,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();
@@ -116,7 +116,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')
@@ -132,8 +132,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();
@@ -141,7 +141,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/ConfigUITest.php b/src/Tests/ConfigUITest.php
index 3358439..e66c960 100644
--- a/src/Tests/ConfigUITest.php
+++ b/src/Tests/ConfigUITest.php
@@ -5,14 +5,14 @@ namespace Drupal\entity_browser\Tests;
 use Drupal\entity_browser\Plugin\EntityBrowser\Display\IFrame;
 use Drupal\entity_browser\Plugin\EntityBrowser\SelectionDisplay\NoDisplay;
 use Drupal\entity_browser\Plugin\EntityBrowser\WidgetSelector\Tabs;
-use Drupal\simpletest\WebTestBase;
+use Drupal\FunctionalJavascriptTests\WebDriverTestBase;
 
 /**
  * Tests the entity browser config UI.
  *
  * @group entity_browser
  */
-class ConfigUITest extends WebTestBase {
+class ConfigUITest extends WebDriverTestBase {
 
   /**
    * The test user.
@@ -47,19 +47,19 @@ class ConfigUITest extends WebTestBase {
     $this->container->get('module_installer')->install(['token']);
 
     $this->drupalGet('/admin/config/content/entity_browser');
-    $this->assertResponse(403, "Anonymous user can't access entity browser listing page.");
+    $this->assertSession()->responseNotContains('Add Entity browser');
     $this->drupalGet('/admin/config/content/entity_browser/add');
-    $this->assertResponse(403, "Anonymous user can't access entity browser add form.");
+    $this->assertSession()->responseNotContains('Entity Browser creation instructions');
 
     // Listing is empty.
     $this->drupalLogin($this->adminUser);
     $this->drupalGet('/admin/config/content/entity_browser');
-    $this->assertResponse(200, 'Admin user is able to navigate to the entity browser listing page.');
-    $this->assertText('There are no entity browser entities yet.', 'Entity browsers table is empty.');
+    $this->assertSession()->responseNotContains('Add Entity browser');
+    $this->assertSession()->responseContains('There are no entity browser entities yet.', 'Entity browsers table is empty.');
 
     // Add page.
     $this->clickLink('Add Entity browser');
-    $this->assertUrl('/admin/config/content/entity_browser/add');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/add');
     $edit = [
       'label' => 'Test entity browser',
       'id' => 'test_entity_browser',
@@ -67,76 +67,86 @@ class ConfigUITest extends WebTestBase {
       'widget_selector' => 'tabs',
       'selection_display' => 'no_display',
     ];
-    $this->drupalPostForm(NULL, $edit, 'Next');
+    $this->submitForm($edit, 'Next');
 
     // Display configuration step.
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/display', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/display', ['query' => ['js' => 'nojs']]);
     $edit = [
       'width' => 100,
       'height' => 100,
       'link_text' => 'All animals are created equal',
       'auto_open' => TRUE,
     ];
-    $this->drupalPostForm(NULL, $edit, 'Next');
+    $this->submitForm($edit, 'Next');
 
     // Widget selector step.
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/widget_selector', ['query' => ['js' => 'nojs']]);
-    $this->assertText('This plugin has no configuration options.');
-    $this->drupalPostForm(NULL, [], 'Next');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/widget_selector', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->responseContains('This plugin has no configuration options.');
+    $this->submitForm([], 'Next');
 
     // Selection display step.
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/selection_display', ['query' => ['js' => 'nojs']]);
-    $this->assertText('This plugin has no configuration options.');
-    $this->drupalPostForm(NULL, [], 'Previous');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/selection_display', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->responseContains('This plugin has no configuration options.');
+    $this->submitForm([], 'Previous');
 
     // Widget selector step again.
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/widget_selector', ['query' => ['js' => 'nojs']]);
-    $this->assertText('This plugin has no configuration options.');
-    $this->drupalPostForm(NULL, [], 'Next');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/widget_selector', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->responseContains('This plugin has no configuration options.');
+    $this->submitForm([], 'Next');
 
     // Selection display step.
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/selection_display', ['query' => ['js' => 'nojs']]);
-    $this->assertText('This plugin has no configuration options.');
-    $this->drupalPostForm(NULL, [], 'Next');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/selection_display', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->responseContains('This plugin has no configuration options.');
+    $this->submitForm([], 'Next');
 
     // Widgets step.
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/widgets', ['query' => ['js' => 'nojs']]);
-    $this->assertText('The available plugins are:');
-    $this->assertText("Upload: Adds an upload field browser's widget.");
-    $this->assertText("View: Uses a view to provide entity listing in a browser's widget.");
-    $this->assertText("Entity form: Provides entity form widget.");
-    $this->drupalPostAjaxForm(NULL, ['widget' => 'upload'], 'widget');
-    $this->assertText('Label (Upload)');
-    $this->assertText('You can use tokens in the upload location.');
-    $this->assertLink('Browse available tokens.');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/widgets', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->responseContains('The available plugins are:');
+    $this->assertSession()->responseContains("<b>Upload:</b> Adds an upload field browser's widget.");
+    $this->assertSession()->responseContains("<b>View:</b> Uses a view to provide entity listing in a browser's widget.");
+    $this->assertSession()->responseContains("<b>Entity form:</b> Provides entity form widget.");
+    $this->assertSession()->selectExists('widget')->setValue('upload');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->responseContains('Label (Upload)');
+    $this->assertSession()->responseContains('You can use tokens in the upload location.');
+    $this->assertSession()->linkExists('Browse available tokens.');
 
     // Make sure that removing of widgets works.
-    $this->drupalPostAjaxForm(NULL, ['widget' => 'view'], 'widget');
-    $this->assertText('Label (View)');
-    $this->assertText('View : View display', 'View selection dropdown label found.');
-    $this->assertRaw('- Select a view -', 'Empty option appears in the view selection dropdown.');
-    $this->assertText('Submit button text', 'Widget submit button text element found.');
-    $this->assertFieldByXPath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-submit-text")]', 'Select entities', 'Widget submit button text element found.');
+    $this->assertSession()->selectExists('widget')->setValue('view');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->responseContains('Label (View)');
+    $this->assertSession()->responseContains('View : View display', 'View selection dropdown label found.');
+    $this->assertSession()->responseContains('- Select a view -', 'Empty option appears in the view selection dropdown.');
+    $this->assertSession()->responseContains('Submit button text', 'Widget submit button text element found.');
+
+    $fields = $this->xpath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-submit-text")]');
+    $this->assertFieldsByValue($fields, 'Select entities', 'Widget submit button text element found.');
+
     $delete_buttons = $this->xpath("//input[@value='Delete']");
     $delete_button_name = (string) $delete_buttons[1]->attributes()['name'];
-    $this->drupalPostAjaxForm(NULL, [], [$delete_button_name => 'Delete']);
-    $this->assertNoText('View : View display', 'View widget was removed.');
-    $this->assertNoRaw('- Select a view -', 'View widget was removed.');
-    $this->assertEqual(count($this->xpath("//input[@value='Delete']")), 1, 'Only one delete button appears on the page.');
+    $this->submitForm([], [$delete_button_name => 'Delete']);
+    $this->assertSession()->responseNotContains('View : View display', 'View widget was removed.');
+    $this->assertSession()->responseNotContains('- Select a view -', 'View widget was removed.');
+    $this->assertEquals(count($this->xpath("//input[@value='Delete']")), 1, 'Only one delete button appears on the page.');
 
     // Make sure the "Entity form" widget has all available config elements.
-    $this->drupalPostAjaxForm(NULL, ['widget' => 'entity_form'], 'widget');
-    $this->assertText('Label (Entity form)');
-    $this->assertText('Entity type', 'Entity type select found on IEF widget.');
-    $this->assertText('Bundle', 'Bundle select found on IEF widget.');
-    $this->assertText('Form mode', 'Form mode select found on IEF widget.');
-    $this->assertFieldByXPath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-submit-text")]', 'Save entity', 'Widget submit button text element found.');
+    $this->assertSession()->selectExists('widget')->setValue('entity_form');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->assertSession()->responseContains('Label (Entity form)');
+    $this->assertSession()->responseContains('Entity type', 'Entity type select found on IEF widget.');
+    $this->assertSession()->responseContains('Bundle', 'Bundle select found on IEF widget.');
+    $this->assertSession()->responseContains('Form mode', 'Form mode select found on IEF widget.');
+
+    $fields = $this->xpath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-submit-text")]');
+
+    $this->assertFieldsByValue($fields, 'Save entity', 'Widget submit button text element found.');
+
     $entity_type_element = $this->xpath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-entity-type")]');
     $entity_type_name = (string) $entity_type_element[0]['name'];
     $edit = [
       $entity_type_name => 'user',
     ];
-    $commands = $this->drupalPostAjaxForm(NULL, $edit, $entity_type_name);
+    $commands = $this->submitForm($edit, $entity_type_name);
     // WebTestBase::drupalProcessAjaxResponse() won't correctly execute our ajax
     // commands so we have to do it manually. Code below is based on the logic
     // in that function.
@@ -160,7 +170,7 @@ class ConfigUITest extends WebTestBase {
     $form_mode_element = $this->xpath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-form-mode-form-select")]');
     $form_mode_id = (string) $form_mode_element[0]['id'];
     $form_mode_name = (string) $form_mode_element[0]['name'];
-    $this->assertOption($form_mode_id, 'register', 'A non-default form mode is correctly available to be chosen.');
+    $this->assertSession()->optionExists($form_mode_id, 'register', 'A non-default form mode is correctly available to be chosen.');
     $bundle_element = $this->xpath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-bundle-select")]');
     $bundle_name = (string) $bundle_element[0]['name'];
     $submit_text_element = $this->xpath('//*[starts-with(@data-drupal-selector, "edit-table-") and contains(@data-drupal-selector, "-form-submit-text")]');
@@ -171,20 +181,20 @@ class ConfigUITest extends WebTestBase {
       $form_mode_name => 'register',
       $submit_text_name => 'But some are more equal than others',
     ];
-    $this->drupalPostForm(NULL, $edit, 'Finish');
+    $this->submitForm($edit, 'Finish');
 
     // Back on listing page.
-    $this->assertUrl('/admin/config/content/entity_browser');
-    $this->assertText('Test entity browser', 'Entity browser label found on the listing page');
-    $this->assertText('test_entity_browser', 'Entity browser ID found on the listing page.');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser');
+    $this->assertSession()->responseContains('Test entity browser', 'Entity browser label found on the listing page');
+    $this->assertSession()->responseContains('test_entity_browser', 'Entity browser ID found on the listing page.');
 
     // Check structure of entity browser object.
     /** @var \Drupal\entity_browser\EntityBrowserInterface $loaded_entity_browser */
     $loaded_entity_browser = $this->container->get('entity_type.manager')
       ->getStorage('entity_browser')
       ->load('test_entity_browser');
-    $this->assertEqual('test_entity_browser', $loaded_entity_browser->id(), 'Entity browser ID was correctly saved.');
-    $this->assertEqual('Test entity browser', $loaded_entity_browser->label(), 'Entity browser label was correctly saved.');
+    $this->assertEquals('test_entity_browser', $loaded_entity_browser->id(), 'Entity browser ID was correctly saved.');
+    $this->assertEquals('Test entity browser', $loaded_entity_browser->label(), 'Entity browser label was correctly saved.');
     $this->assertTrue($loaded_entity_browser->getDisplay() instanceof IFrame, 'Entity browser display was correctly saved.');
     $expected = [
       'width' => '100',
@@ -192,11 +202,11 @@ class ConfigUITest extends WebTestBase {
       'link_text' => 'All animals are created equal',
       'auto_open' => TRUE,
     ];
-    $this->assertEqual($expected, $loaded_entity_browser->getDisplay()->getConfiguration(), 'Entity browser display configuration was correctly saved.');
+    $this->assertEquals($expected, $loaded_entity_browser->getDisplay()->getConfiguration(), 'Entity browser display configuration was correctly saved.');
     $this->assertTrue($loaded_entity_browser->getSelectionDisplay() instanceof NoDisplay, 'Entity browser selection display was correctly saved.');
-    $this->assertEqual([], $loaded_entity_browser->getSelectionDisplay()->getConfiguration(), 'Entity browser selection display configuration was correctly saved.');
-    $this->assertEqual($loaded_entity_browser->getWidgetSelector() instanceof Tabs, 'Entity browser widget selector was correctly saved.');
-    $this->assertEqual([], $loaded_entity_browser->getWidgetSelector()->getConfiguration(), 'Entity browser widget selector configuration was correctly saved.');
+    $this->assertEquals([], $loaded_entity_browser->getSelectionDisplay()->getConfiguration(), 'Entity browser selection display configuration was correctly saved.');
+    $this->assertEquals($loaded_entity_browser->getWidgetSelector() instanceof Tabs, 'Entity browser widget selector was correctly saved.');
+    $this->assertEquals([], $loaded_entity_browser->getWidgetSelector()->getConfiguration(), 'Entity browser widget selector configuration was correctly saved.');
 
     $widgets = $loaded_entity_browser->getWidgets();
     $instance_ids = $widgets->getInstanceIds();
@@ -204,78 +214,83 @@ class ConfigUITest extends WebTestBase {
     $second_uuid = next($instance_ids);
     /** @var \Drupal\entity_browser\WidgetInterface $widget */
     $widget = $widgets->get($first_uuid);
-    $this->assertEqual('upload', $widget->id(), 'Entity browser widget was correctly saved.');
-    $this->assertEqual($first_uuid, $widget->uuid(), 'Entity browser widget uuid was correctly saved.');
+    $this->assertEquals('upload', $widget->id(), 'Entity browser widget was correctly saved.');
+    $this->assertEquals($first_uuid, $widget->uuid(), 'Entity browser widget uuid was correctly saved.');
     $configuration = $widget->getConfiguration()['settings'];
-    $this->assertEqual([
+    $this->assertEquals([
       'upload_location' => 'public://',
       'multiple' => TRUE,
       'submit_text' => 'Select files',
       'extensions' => 'jpg jpeg gif png txt doc xls pdf ppt pps odt ods odp',
     ], $configuration, 'Entity browser widget configuration was correctly saved.');
-    $this->assertEqual(1, $widget->getWeight(), 'Entity browser widget weight was correctly saved.');
+    $this->assertEquals(1, $widget->getWeight(), 'Entity browser widget weight was correctly saved.');
     $widget = $widgets->get($second_uuid);
-    $this->assertEqual('entity_form', $widget->id(), 'Entity browser widget was correctly saved.');
-    $this->assertEqual($second_uuid, $widget->uuid(), 'Entity browser widget uuid was correctly saved.');
+    $this->assertEquals('entity_form', $widget->id(), 'Entity browser widget was correctly saved.');
+    $this->assertEquals($second_uuid, $widget->uuid(), 'Entity browser widget uuid was correctly saved.');
     $configuration = $widget->getConfiguration()['settings'];
-    $this->assertEqual([
+    $this->assertEquals([
       'entity_type' => 'user',
       'bundle' => 'user',
       'form_mode' => 'register',
       'submit_text' => 'But some are more equal than others',
     ], $configuration, 'Entity browser widget configuration was correctly saved.');
-    $this->assertEqual(2, $widget->getWeight(), 'Entity browser widget weight was correctly saved.');
+    $this->assertEquals(2, $widget->getWeight(), 'Entity browser widget weight was correctly saved.');
 
     // Navigate to edit.
     $this->clickLink('Edit');
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser');
-    $this->assertFieldById('edit-label', 'Test entity browser', 'Correct label found.');
-    $this->assertText('test_entity_browser', 'Correct id found.');
-    $this->assertOptionSelected('edit-display', 'iframe', 'Correct display selected.');
-    $this->assertOptionSelected('edit-widget-selector', 'tabs', 'Correct widget selector selected.');
-    $this->assertOptionSelected('edit-selection-display', 'no_display', 'Correct selection display selected.');
-
-    $this->drupalPostForm(NULL, [], 'Next');
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/display', ['query' => ['js' => 'nojs']]);
-    $this->assertFieldById('edit-width', '100', 'Correct value for width found.');
-    $this->assertFieldById('edit-height', '100', 'Correct value for height found.');
-    $this->assertFieldById('edit-link-text', 'All animals are created equal', 'Correct value for link text found.');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser');
+    $this->assertSession()->fieldValueEquals('label', 'Test entity browser');
+    $this->assertSession()->responseContains('test_entity_browser', 'Correct id found.');
+
+    $option_field = $this->assertSession()->optionExists('edit-display', 'iframe');
+    $this->assertTrue($option_field->hasAttribute('selected'), 'Correct display selected.');
+    $option_field = $this->assertSession()->optionExists('edit-widget-selector', 'tabs');
+    $this->assertTrue($option_field->hasAttribute('selected'), 'Correct widget selector selected.');
+    $option_field = $this->assertSession()->optionExists('edit-selection-display', 'no_display');
+    $this->assertTrue($option_field->hasAttribute('selected'), 'Correct selection display selected.');
+
+    $this->submitForm([], 'Next');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/display', ['query' => ['js' => 'nojs']]);
+    $this->assertSession()->fieldValueEquals('width', '100');
+    $this->assertSession()->fieldValueEquals('height', '100');
+    $this->assertSession()->fieldValueEquals('link_text', 'All animals are created equal');
     $this->assertFieldChecked('edit-auto-open', 'Auto open is enabled.');
 
-    $this->drupalPostForm(NULL, [], 'Next');
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/widget_selector', ['query' => ['js' => 'nojs']]);
-
-    $this->drupalPostForm(NULL, [], 'Next');
-    $this->assertUrl('/admin/config/content/entity_browser/test_entity_browser/selection_display', ['query' => ['js' => 'nojs']]);
-
-    $this->drupalPostForm(NULL, [], 'Next');
-    $this->assertFieldById('edit-table-' . $first_uuid . '-label', 'upload', 'Correct value for widget label found.');
-    $this->assertFieldChecked('edit-table-' . $first_uuid . '-form-multiple', 'Accept multiple files option is enabled by default.');
-    $this->assertText('Multiple uploads will only be accepted if the source field allows more than one value.');
-    $this->assertFieldById('edit-table-' . $first_uuid . '-form-upload-location', 'public://', 'Correct value for upload location found.');
-    $this->assertFieldByXPath("//input[@data-drupal-selector='edit-table-" . $first_uuid . "-form-submit-text']", 'Select files', 'Correct value for submit text found.');
-    $this->assertFieldById('edit-table-' . $second_uuid . '-label', 'entity_form', 'Correct value for widget label found.');
-    $this->assertOptionSelectedWithDrupalSelector('edit-table-' . $second_uuid . '-form-entity-type', 'user', 'Correct value for entity type found.');
-    $this->assertOptionSelectedWithDrupalSelector('edit-table-' . $second_uuid . '-form-bundle-select', 'user', 'Correct value for bundle found.');
-    $this->assertOptionSelectedWithDrupalSelector('edit-table-' . $second_uuid . '-form-form-mode-form-select', 'register', 'Correct value for form modes found.');
-    $this->assertFieldByXPath("//input[@data-drupal-selector='edit-table-" . $second_uuid . "-form-submit-text']", 'But some are more equal than others', 'Correct value for submit text found.');
-
-    $this->drupalPostForm(NULL, ['table[' . $first_uuid . '][form][multiple]' => FALSE], 'Finish');
+    $this->submitForm([], 'Next');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/widget_selector', ['query' => ['js' => 'nojs']]);
+
+    $this->submitForm([], 'Next');
+    $this->assertSession()->addressEquals('/admin/config/content/entity_browser/test_entity_browser/selection_display', ['query' => ['js' => 'nojs']]);
+
+    $this->submitForm([], 'Next');
+    $this->assertSession()->fieldValueEquals("table[$first_uuid][label]", "upload");
+    $this->assertSession()->checkboxChecked("table[$first_uuid][form][multiple]");
+    $this->assertSession()->responseContains('Multiple uploads will only be accepted if the source field allows more than one value.');
+    $this->assertSession()->fieldValueEquals("table[$first_uuid][form][upload_location]", 'public://');
+    $this->assertSession()->fieldValueEquals("table[$first_uuid][form][submit_text]", 'Select files');
+
+    $this->assertSession()->fieldValueEquals("table[$second_uuid][label]", 'entity_form');
+
+    $this->assertSession()->fieldValueEquals("table[$second_uuid][form][entity_type]", 'user', 'Correct value for entity type found.');
+    $this->fieldValueEquals("table[$second_uuid][form][bundle][select]", 'user', 'Correct value for bundle found.');
+    $this->fieldValueEquals("table[$second_uuid][form][form_mode][form_select]", 'register', 'Correct value for form modes found.');
+    $this->assertSession()->fieldValueEquals("table[$second_uuid][form][submit_text]", 'But some are more equal than others');
+    $this->submitForm(['table[' . $first_uuid . '][form][multiple]' => FALSE], 'Finish');
     $this->drupalGet('/admin/config/content/entity_browser/test_entity_browser/widgets');
-    $this->assertNoFieldChecked('edit-table-' . $first_uuid . '-form-multiple', 'Accept multiple files option is disabled.');
+    $this->assertSession()->checkboxNotChecked("table[$first_uuid][form][multiple]", 'Accept multiple files option is disabled.');
 
     $this->drupalLogout();
     $this->drupalGet('/admin/config/content/entity_browser/test_entity_browser');
-    $this->assertResponse(403, "Anonymous user can't access entity browser edit form.");
+    $this->assertSession()->statusCodeEquals(403, "Anonymous user can't access entity browser edit form.");
 
     $this->drupalLogin($this->adminUser);
     $this->drupalGet('/admin/config/content/entity_browser');
     $this->clickLink('Delete');
-    $this->assertText('This action cannot be undone.', 'Delete question found.');
-    $this->drupalPostForm(NULL, [], 'Delete Entity Browser');
+    $this->assertSession()->responseContains('This action cannot be undone.', 'Delete question found.');
+    $this->submitForm([], 'Delete Entity Browser');
 
-    $this->assertText('Entity browser Test entity browser was deleted.', 'Confirmation message found.');
-    $this->assertText('There are no entity browser entities yet.', 'Entity browsers table is empty.');
+    $this->assertSession()->responseContains('Entity browser Test entity browser was deleted.', 'Confirmation message found.');
+    $this->assertSession()->responseContains('There are no entity browser entities yet.', 'Entity browsers table is empty.');
     $this->drupalLogout();
   }
 
diff --git a/src/Tests/EntityBrowserUITest.php b/src/Tests/EntityBrowserUITest.php
index 5e45754..351337c 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.
@@ -35,14 +40,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, 'Finish');
+    $this->submitForm($edit, 'Finish');
     $this->drupalGet('/entity-browser/iframe/test_entity_browser_iframe');
-    $this->assertRaw('Different');
+    $this->assertSession()->responseContains('Different');
   }
 
   /**
@@ -56,16 +61,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/src/Tests/MultistepDisplayTest.php b/src/Tests/MultistepDisplayTest.php
index 1d91a5d..7ec1025 100644
--- a/src/Tests/MultistepDisplayTest.php
+++ b/src/Tests/MultistepDisplayTest.php
@@ -2,14 +2,20 @@
 
 namespace Drupal\entity_browser\Tests;
 
-use Drupal\simpletest\WebTestBase;
+use Drupal\FunctionalJavascriptTests\WebDriverTestBase;
+use Drupal\Tests\TestFileCreationTrait;
 
 /**
  * Tests the multistep display selection display.
  *
  * @group entity_browser
  */
-class MultistepDisplayTest extends WebTestBase {
+class MultistepDisplayTest extends WebDriverTestBase {
+
+  use TestFileCreationTrait {
+    getTestFiles as drupalGetTestFiles;
+    compareFiles as drupalCompareFiles;
+  }
 
   /**
    * Modules to enable.
@@ -43,23 +49,24 @@ class MultistepDisplayTest extends WebTestBase {
       'widget_selector' => 'tabs',
       'selection_display' => 'multi_step_display',
     ];
-    $this->drupalPostForm(NULL, $edit, 'Next');
-    $this->drupalPostForm(NULL, [], 'Next');
-    $this->drupalPostForm(NULL, [], 'Next');
+    $this->submitForm($edit, 'Next');
+    $this->submitForm([], 'Next');
+    $this->submitForm([], 'Next');
 
-    $this->assertText('Selection display', 'Trail is shown.');
-    $this->assertText('Select button text', 'Title is correct.');
-    $this->assertText('Text to display on the entity browser select button.', 'Description is correct.');
-    $this->assertRaw('Use selected', 'Default text is correct.');
+    $this->assertSession()->responseContains('Selection display', 'Trail is shown.');
+    $this->assertSession()->responseContains('Select button text', 'Title is correct.');
+    $this->assertSession()->responseContains('Text to display on the entity browser select button.', 'Description is correct.');
+    $this->assertSession()->responseContains('Use selected', 'Default text is correct.');
     $edit = [
       'entity_type' => 'file',
       'display' => 'label',
       'selection_hidden' => 0,
       'select_text' => 'Use blah selected',
     ];
-    $this->drupalPostForm(NULL, $edit, 'Next');
-    $this->drupalPostAjaxForm(NULL, ['widget' => 'upload'], 'widget');
-    $this->drupalPostForm(NULL, [], 'Finish');
+    $this->submitForm($edit, 'Next');
+    $this->assertSession()->selectExists('widget')->setValue('upload');
+    $this->assertSession()->assertWaitOnAjaxRequest();
+    $this->submitForm([], 'Finish');
 
     $account = $this->drupalCreateUser([
       'access test_entity_browser entity browser pages',
@@ -67,14 +74,14 @@ class MultistepDisplayTest extends WebTestBase {
     $this->drupalLogin($account);
     // Go to the entity browser iframe link.
     $this->drupalGet('/entity-browser/iframe/test_entity_browser');
-    $this->assertNoRaw('Use blah selected');
+    $this->assertSession()->responseNotContains('Use blah selected');
 
     $image = current($this->drupalGetTestFiles('image'));
     $edit = [
       'files[upload][]' => $this->container->get('file_system')->realpath($image->uri),
     ];
-    $this->drupalPostForm(NULL, $edit, 'Select files');
-    $this->assertRaw('Use blah selected', 'Select button is displayed if something is selected.');
+    $this->submitForm($edit, 'Select files');
+    $this->assertSession()->responseContains('Use blah selected', 'Select button is displayed if something is selected.');
   }
 
 }
diff --git a/src/Wizard/EntityBrowserWizard.php b/src/Wizard/EntityBrowserWizard.php
index 4174bcb..c6b3362 100644
--- a/src/Wizard/EntityBrowserWizard.php
+++ b/src/Wizard/EntityBrowserWizard.php
@@ -3,16 +3,16 @@
 namespace Drupal\entity_browser\Wizard;
 
 use Drupal\Core\DependencyInjection\ClassResolverInterface;
-use Drupal\Core\Entity\EntityManagerInterface;
+use Drupal\Core\Entity\EntityTypeManagerInterface;
 use Drupal\Core\Form\FormBuilderInterface;
 use Drupal\Core\Routing\RouteMatchInterface;
+use Drupal\Core\TempStore\SharedTempStoreFactory;
 use Drupal\ctools\Wizard\EntityFormWizardBase;
 use Drupal\entity_browser\Form\DisplayConfig;
 use Drupal\entity_browser\Form\GeneralInfoConfig;
 use Drupal\entity_browser\Form\SelectionDisplayConfig;
 use Drupal\entity_browser\Form\WidgetsConfig;
 use Drupal\entity_browser\Form\WidgetSelectorConfig;
-use Drupal\user\SharedTempStoreFactory;
 use Symfony\Component\EventDispatcher\EventDispatcherInterface;
 
 /**
@@ -21,7 +21,7 @@ use Symfony\Component\EventDispatcher\EventDispatcherInterface;
 class EntityBrowserWizard extends EntityFormWizardBase {
 
   /**
-   * @param \Drupal\user\SharedTempStoreFactory $tempstore
+   * @param \Drupal\Core\TempStore\SharedTempStoreFactory $tempstore
    *   Tempstore Factory for keeping track of values in each step of the
    *   wizard.
    * @param \Drupal\Core\Form\FormBuilderInterface $builder
@@ -30,8 +30,10 @@ class EntityBrowserWizard extends EntityFormWizardBase {
    *   The class resolver.
    * @param \Symfony\Component\EventDispatcher\EventDispatcherInterface $event_dispatcher
    *   The event dispatcher.
-   * @param \Drupal\Core\Entity\EntityManagerInterface $entity_manager
+   * @param \Drupal\Core\Entity\EntityTypeManagerInterface $entity_type_manager
    *   The entity manager.
+   * @param \Drupal\Core\Routing\RouteMatchInterface $route_match
+   *   The route match
    * @param $tempstore_id
    *   The shared temp store factory collection name.
    * @param null $machine_name
@@ -39,8 +41,8 @@ class EntityBrowserWizard extends EntityFormWizardBase {
    * @param null $step
    *   The current active step of the wizard.
    */
-  public function __construct(SharedTempStoreFactory $tempstore, FormBuilderInterface $builder, ClassResolverInterface $class_resolver, EventDispatcherInterface $event_dispatcher, EntityManagerInterface $entity_manager, RouteMatchInterface $route_match, $tempstore_id, $entity_browser = NULL, $step = 'general') {
-    parent::__construct($tempstore, $builder, $class_resolver, $event_dispatcher, $entity_manager, $route_match, $tempstore_id, $entity_browser, $step);
+  public function __construct(SharedTempStoreFactory $tempstore, FormBuilderInterface $builder, ClassResolverInterface $class_resolver, EventDispatcherInterface $event_dispatcher, EntityTypeManagerInterface $entity_type_manager, RouteMatchInterface $route_match, $tempstore_id, $entity_browser = NULL, $step = 'general') {
+    parent::__construct($tempstore, $builder, $class_resolver, $event_dispatcher, $entity_type_manager, $route_match, $tempstore_id, $entity_browser, $step);
   }
 
   /**
diff --git a/src/Tests/EntityBrowserUpdateHookTest.php b/tests/src/Functional/EntityBrowserUpdateHookTest.php
similarity index 91%
rename from src/Tests/EntityBrowserUpdateHookTest.php
rename to tests/src/Functional/EntityBrowserUpdateHookTest.php
index e33ac14..d644c8d 100644
--- a/src/Tests/EntityBrowserUpdateHookTest.php
+++ b/tests/src/Functional/EntityBrowserUpdateHookTest.php
@@ -1,8 +1,9 @@
 <?php
 
-namespace Drupal\entity_browser\Tests;
+namespace Drupal\Tests\entity_browser\Functional;
 
-use Drupal\system\Tests\Update\UpdatePathTestBase;
+
+use Drupal\FunctionalTests\Update\UpdatePathTestBase;
 
 /**
  * Tests the update hooks in entity_browser module.
@@ -17,7 +18,7 @@ class EntityBrowserUpdateHookTest extends UpdatePathTestBase {
   protected function setDatabaseDumpFiles() {
     $this->databaseDumpFiles = [
       DRUPAL_ROOT . '/core/modules/system/tests/fixtures/update/drupal-8.bare.standard.php.gz',
-      __DIR__ . '/../../tests/fixtures/update/entity_browser.update-hook-test.php',
+      __DIR__ . '/../../fixtures/update/entity_browser.update-hook-test.php',
     ];
   }
 
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 568a595..585660f 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.
@@ -211,7 +217,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 a9f761c..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,22 +188,28 @@ 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()->pageTextContains('Only files with the following extensions are allowed: jpg');
-    $this->assertSession()->pageTextContains('The specified file druplicon.png could not be uploaded');
+    $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.');
+    }
+    else {
+      $this->assertSession()->pageTextContains('Only files with the following extensions are allowed: jpg');
+      $this->assertSession()->pageTextContains('The specified file druplicon.png could not be uploaded');
+    }
     // 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 a3886b3..02ba032 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}
@@ -90,7 +90,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();
 
@@ -102,22 +102,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.
@@ -125,10 +125,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).
@@ -142,10 +142,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.
@@ -159,17 +159,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).
@@ -189,14 +189,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).
@@ -213,17 +213,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).
@@ -247,13 +247,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 f3d5036..b6400ac 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}
@@ -52,30 +52,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';
@@ -85,42 +85,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');
@@ -135,7 +136,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 018f17d..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;
 
 /**
@@ -128,7 +127,7 @@ class EntityBrowserTest extends KernelTestBase {
         $this->fail('An entity browser without required ' . $plugin_type . ' created with no exception thrown.');
       }
       catch (PluginException $e) {
-        $this->assertEquals('The "" plugin does not exist.', $e->getMessage(), 'An exception was thrown when an entity_browser was created without a ' . $plugin_type . ' plugin.');
+        $this->assertContains('The "" plugin does not exist.', $e->getMessage(), 'An exception was thrown when an entity_browser was created without a ' . $plugin_type . ' plugin.');
       }
     }
 
