diff --git a/src/Controller/ActionLinkController.php b/src/Controller/ActionLinkController.php
index b13c476..ce7dd84 100644
--- a/src/Controller/ActionLinkController.php
+++ b/src/Controller/ActionLinkController.php
@@ -89,7 +89,7 @@ class ActionLinkController extends ControllerBase implements ContainerInjectionI
       // link for the existing state of the flag.
     }
 
-    return $this->generateResponse($flag, $entity, $request, $flag->getFlagMessage());
+    return $this->generateResponse($flag, $entity, $request, $flag->getMessage('flag'));
   }
 
   /**
@@ -119,7 +119,7 @@ class ActionLinkController extends ControllerBase implements ContainerInjectionI
       // link for the existing state of the flag.
     }
 
-    return $this->generateResponse($flag, $entity, $request, $flag->getUnflagMessage());
+    return $this->generateResponse($flag, $entity, $request, $flag->getMessage('unflag'));
   }
 
   /**
diff --git a/tests/src/FunctionalJavascript/AjaxLinkTest.php b/tests/src/FunctionalJavascript/AjaxLinkTest.php
new file mode 100644
index 0000000..50b9364
--- /dev/null
+++ b/tests/src/FunctionalJavascript/AjaxLinkTest.php
@@ -0,0 +1,124 @@
+<?php
+
+namespace Drupal\Tests\flag\FunctionalJavascript;
+
+use Drupal\flag\Tests\FlagCreateTrait;
+use Drupal\FunctionalJavascriptTests\JavascriptTestBase;
+
+/**
+ * Javascript test for AjaxLinks.
+ *
+ * When a user clicks on an AJAX link a salvo of AJAX commands is issued in
+ * response which update the DOM with a new link and a short lived message.
+ *
+ * @group flag
+ */
+class AjaxLinkTest extends JavascriptTestBase {
+
+  use FlagCreateTrait;
+
+  /**
+   * {@inheritdoc}
+   */
+  public static $modules = ['flag', 'node', 'user'];
+
+  /**
+   * Flag to test with.
+   *
+   * @var \Drupal\flag\FlagInterface
+   */
+  protected $flag;
+
+  /**
+   * The flag service.
+   *
+   * @var \Drupal\flag\FlagServiceInterface
+   */
+  protected $flagService;
+
+  /**
+   * Test node.
+   *
+   * @var \Drupal\node\NodeInterface
+   */
+  protected $node;
+
+  /**
+   * Admin user.
+   *
+   * @var \Drupal\user\UserInterface
+   */
+  protected $admin;
+
+  /**
+   * Normal user.
+   *
+   * @var \Drupal\user\UserInterface
+   */
+  protected $webUser;
+
+  /**
+   * {@inheritdoc}
+   */
+  protected function setUp() {
+    parent::setUp();
+
+    // A test flag.
+    $this->flag = $this->createFlag('node', ['article'], 'ajax_link');
+    $this->flagService = $this->container->get('flag');
+
+    // A article to test with.
+    $this->admin = $this->createUser([], NULL, TRUE);
+    $type = $this->createContentType(['type' => 'article']);
+    $this->node = $this->createNode([
+      'type' => 'article',
+      'uid' => $this->admin->id(),
+    ]);
+
+    $this->webUser = $this->createUser(array_keys($this->flag->actionPermissions()));
+    $this->drupalLogin($this->webUser);
+  }
+
+  /**
+   * Tests DOM update.
+   *
+   * Click on flag and then unflag links verifying that the link cycles as
+   * expected and flag flash message functions.
+   */
+  public function testDomUpdate() {
+    // Get Page.
+    $node_path = '/node/' . $this->node->id();
+    $this->drupalGet($node_path);
+    $session = $this->getSession();
+
+    // Verify initially flag link is on the page.
+    $page = $session->getPage();
+    $this->createScreenshot('/Users/martin/sites/drupal/screenshot.jpg');
+    $flag_link = $page->findLink($this->flag->getShortText('flag'));
+    $this->assertTrue($flag_link->isVisible(), 'flag link exists.');
+
+    $flag_link->click();
+
+    // Verify flags message appears.
+    $flag_message = $this->flag->getMessage('flag');
+    $p_flash = $this->assertSession()->waitForElementVisible('css', "p.js-flag-message");
+    $this->assertEquals($p_flash->getText(), $flag_message, 'DOM update: The flag message is flashed.');
+
+    // Verify new link.
+    $unflag_link = $session->getPage()->findLink($this->flag->getShortText('unflag'));
+    $this->assertTrue($unflag_link->isVisible(), 'unflag link exists.');
+
+    $unflag_link->click();
+
+    // Verfy unflag message appears.
+    $unflag_message = $this->flag->getMessage('unflag');
+    $p_flash2 = $this->assertSession()->waitForElementVisible('css', "p.js-flag-message");
+    $this->assertEquals($p_flash2->getText(), $unflag_message, 'DOM update: The unflag message is flashed.');
+
+    // Verify the cycle completes and flag returns.
+    $flag_link2 = $session->getPage()->findLink($this->flag->getShortText('flag'));
+    $this->assertTrue($flag_link2->isVisible(), 'flag cycle return to start.');
+
+  }
+
+}