diff --git a/tests/src/FunctionalJavascript/LinkitDialogCKEditor5Test.php b/tests/src/FunctionalJavascript/LinkitDialogCKEditor5Test.php
index 1e78f45..b2c625b 100644
--- a/tests/src/FunctionalJavascript/LinkitDialogCKEditor5Test.php
+++ b/tests/src/FunctionalJavascript/LinkitDialogCKEditor5Test.php
@@ -21,6 +21,7 @@ use Symfony\Component\Validator\ConstraintViolation;
 class LinkitDialogCKEditor5Test extends WebDriverTestBase {
 
   use ProfileCreationTrait;
+  use CKEditor5TestTrait;
 
   /**
    * {@inheritdoc}
@@ -126,18 +127,13 @@ class LinkitDialogCKEditor5Test extends WebDriverTestBase {
     $entity = EntityTestMul::create(['name' => 'Foo']);
     $entity->save();
 
-    // Go to node creation page.
     $this->drupalGet('node/add/page');
-
-    // Wait until the editor has been loaded.
-    $ckeditor5_loaded = $this->getSession()->wait(5000, "document.querySelectorAll('.ck-toolbar__items').length > 0");
-    $this->assertTrue($ckeditor5_loaded, 'The editor has been loaded.');
-
-    // Click on the drupallink plugin.
-    $page->find('css', '.ck-toolbar__items .ck-button')->click();
+    $this->waitForEditor();
+    $this->pressEditorButton('Link');
 
     // Find the href field.
-    $autocomplete_field = $assert_session->waitForElementVisible('css', '.ck-link-form .ck-input-text');
+    $balloon = $this->assertVisibleBalloon('.ck-link-form');
+    $autocomplete_field = $balloon->find('css', '.ck-input-text');
 
     // Make sure all fields are empty.
     $this->assertEmpty($autocomplete_field->getValue(), 'Autocomplete field is empty.');
@@ -163,8 +159,9 @@ class LinkitDialogCKEditor5Test extends WebDriverTestBase {
 
     // Make sure the linkit field field is populated with the node url.
     $this->assertEquals($entity->toUrl()->toString(), $autocomplete_field->getValue(), 'The href field is populated with the node url.');
-    $page->find('css', '.ck-button-save')->press();
-    $this->assertTrue($assert_session->waitForElementRemoved('css', '.ck-button-save'));
+    $balloon->pressButton('Save');
+    // Assert balloon was closed by pressing its "Save" button.
+    $this->assertFalse($page->find('css', '.ck-balloon-panel')->isVisible());
 
     // Make sure all other attributes are populated.
     $linkit_link = $assert_session->waitForElementVisible('css', '[href][data-entity-type][data-entity-uuid][data-entity-uuid]');
@@ -176,33 +173,20 @@ class LinkitDialogCKEditor5Test extends WebDriverTestBase {
 
     // Open the edit link dialog by moving selection to the link and then click
     // the clicking the editing document.
-    $link = base_path() . 'entity_test_mul/manage/1';
-    $javascript = <<<JS
-(function() {
-  const el = document.querySelector('a[href="$link"]');
-  const range = document.createRange();
-  range.setStart(el.firstChild, 1);
-  range.setEnd(el.firstChild, 1);
-  const sel = window.getSelection();
-  sel.removeAllRanges();
-  sel.addRange(range);
-})();
-JS;
-    $this->getSession()->evaluateScript($javascript);
-    $assert_session->waitForElementVisible('css', '.ck-editor__editable')->click();
-    $link_actions = $assert_session->waitForElementVisible('xpath', '//div[contains(@class, "ck-link-actions")]//button[1]');
-    $this->assertNotNull($link_actions);
-    $link_actions->click();
-    $link_edit_balloon = $assert_session->waitForElementVisible('css', '.ck-link-form');
-    $this->assertNotNull($link_edit_balloon);
-
+    $this->assertFalse($this->getEditorButton('Link')->hasClass('ck-on'));
+    $this->selectTextInsideElement('a');
+    $this->assertTrue($this->getEditorButton('Link')->hasClass('ck-on'));
+    $this->pressEditorButton('Link');
+    $this->assertVisibleBalloon('.ck-link-actions');
+    $edit_button = $this->getBalloonButton('Edit link');
+    $edit_button->click();
+    $link_edit_balloon = $this->assertVisibleBalloon('.ck-link-form');
     $autocomplete_field = $link_edit_balloon->find('css', '.ck-input-text');
-    $this->assertEquals($entity->toUrl()->toString(), $autocomplete_field->getValue(), 'Href field contains the node url when edit.');
+    $this->assertSame($expected_url, $autocomplete_field->getValue());
     // Click to trigger the reset of the the autocomplete status.
     $autocomplete_field->click();
     $autocomplete_field->setValue('http://example.com');
-    $link_edit_balloon->find('css', '.ck-button-save')->click();
-    $assert_session->waitForElementRemoved('css', '.ck-button-save');
+    $link_edit_balloon->pressButton('Save');
 
     $changed_link = $assert_session->waitForElementVisible('css', '.ck-content [href="http://example.com"]');
     $this->assertNotNull($changed_link);