diff --git a/src/Tests/FlagContextualLinksTest.php b/src/Tests/FlagContextualLinksTest.php
deleted file mode 100644
index 4d5f4a4..0000000
--- a/src/Tests/FlagContextualLinksTest.php
+++ /dev/null
@@ -1,167 +0,0 @@
-<?php
-
-namespace Drupal\flag\Tests;
-
-use Drupal\Component\Utility\UrlHelper;
-use Drupal\Component\Serialization\Json;
-use Drupal\Core\Url;
-use Drupal\node\NodeInterface;
-use Drupal\flag\FlagInterface;
-use Drupal\Component\Utility\Html;
-
-/**
- * Test the contextual links with Reload link type.
- *
- * @group flag
- */
-class FlagContextualLinksTest extends FlagTestBase {
-  /**
-   * The flag.
-   *
-   * @var \Drupal\flag\Flaginterface
-   */
-  protected $flag;
-
-  /**
-   * * An authenticated user to test flagging.
-   *
-   * @var \Drupal\user\Entity\User|false
-   */
-  protected $normalUser;
-
-  public function setUp() {
-    parent::setUp();
-
-    // Create a regular user who will be flagging content.
-    $this->normalUser = $this->drupalCreateUser();
-  }
-
-  /**
-   * Create a new flag with the Reload link type and enable contextual links display.
-   */
-  public function testNodeLinks() {
-    $this->doCreateFlag();
-    $this->doFlagNode();
-  }
-
-  /**
-   * Create a node type and flag.
-   */
-  public function doCreateFlag() {
-    $this->drupalLogin($this->adminUser);
-
-    $edit = [
-      'id' => 'test_label_123',
-      'show_as_field' => FALSE,
-      'show_contextual_link' => TRUE,
-    ];
-
-    $this->flag = $this->createFlagWithForm('node', $edit);
-  }
-
-  /**
-   * Create a node and flag it.
-   */
-  public function doFlagNode() {
-    // Grant the flag permissions to the authenticated role.
-    $this->grantFlagPermissions($this->flag);
-
-    $node = $this->drupalCreateNode(['type' => $this->nodeType]);
-
-    // Login as normal user.
-    $this->drupalLogin($this->normalUser);
-
-    // Open node view page that renders Full view mode.
-    $this->drupalGet('node/' . $node->id());
-
-    // Expect Contextual links id for this node to contain flag_keys metadata.
-    $contextual_links_id = 'node:node=' . $node->id() . ':changed=' . $node->getChangedTime() . '&flag_keys=' . $this->flag->id() . '-flag&langcode=en';
-    $this->assertContextualLinkPlaceHolder($contextual_links_id);
-
-    $flag_action_url = $this->getFlagContextualLinkUrl($this->flag, $node, $contextual_links_id, 'flag');
-    $this->assertTrue($flag_action_url, "Flag link found in contextual links.");
-
-    // Emulate click on flag link in Contextual links.
-    $this->drupalGet($flag_action_url);
-
-    // Expect Contextual links id for this node to contain flag_keys metadata with unflag action.
-    $contextual_links_id = 'node:node=' . $node->id() . ':changed=' . $node->getChangedTime() . '&flag_keys=' . $this->flag->id() . '-unflag&langcode=en';
-    $this->assertContextualLinkPlaceHolder($contextual_links_id);
-
-    $flag_action_url = $this->getFlagContextualLinkUrl($this->flag, $node, $contextual_links_id, 'unflag');
-    $this->assertTrue($flag_action_url, "Unflag link found in contextual links.");
-
-    // Login as normal user.
-    $this->drupalLogin($this->drupalCreateUser());
-
-    // Open node view page that renders Full view mode.
-    $this->drupalGet('node/' . $node->id());
-
-    // Expect Contextual links id for this node to contain flag_keys metadata.
-    $contextual_links_id = 'node:node=' . $node->id() . ':changed=' . $node->getChangedTime() . '&flag_keys=' . $this->flag->id() . '-flag&langcode=en';
-    $this->assertContextualLinkPlaceHolder($contextual_links_id);
-
-    $flag_action_url = $this->getFlagContextualLinkUrl($this->flag, $node, $contextual_links_id, 'flag');
-    $this->assertTrue($flag_action_url, "Flag link found in contextual links.");
-  }
-
-  /**
-   * Get server-rendered contextual links for the given contextual link ids.
-   *
-   * @see \Drupal\contextual\Tests\ContextualDynamicContextTest::renderContextualLinks().
-   * @param array $ids
-   *   An array of contextual link ids.
-   * @param string $current_path
-   *   The Drupal path for the page for which the contextual links are rendered.
-   *
-   * @return string
-   *   The response body.
-   */
-  protected function renderContextualLinks($ids, $current_path) {
-    $post = array();
-    for ($i = 0; $i < count($ids); $i++) {
-      $post['ids[' . $i . ']'] = $ids[$i];
-    }
-    return $this->drupalPostWithFormat('contextual/render', 'json', $post, array('query' => array('destination' => $current_path)));
-  }
-
-  /**
-   * Helper method that returns a flag action URL from contextual links markup
-   * by contextual links id.
-   *
-   * @param string $contextual_links_id
-   *   Contextual links placeholder ids string.
-   * @return \Drupal\Core\Url|FALSE
-   *   Returns either a flag action URL object or FALSE in case that flag link
-   *   wasn't found in the contextual links markup.
-   */
-  protected function getFlagContextualLinkUrl(FlagInterface $flag,
-                                              NodeInterface $node,
-                                              $contextual_links_id,
-                                              $flag_action) {
-    // Render contextual links JSON response and check for flag link.
-    $response = $this->renderContextualLinks([$contextual_links_id], 'node');
-    $this->assertResponse(200);
-    $json = Json::decode($response);
-    $flag_url = Url::fromRoute("flag.action_link_{$flag_action}", array(
-      'flag' => $flag->id(),
-      'entity_id' => $node->id(),
-    ));
-
-    $flag_short_text = $flag->getShortText($flag_action);
-
-    $matches = array();
-    $url_pattern =  preg_quote(base_path()) . '(' . preg_quote($flag_url->getInternalPath(), '@') . '[^"]*)';
-    $found = preg_match('@' . preg_quote('<ul class="contextual-links"><li class="flag-test-label-123"><a href="') .
-      $url_pattern . preg_quote('">' . Html::escape($flag_short_text) . '</a></li></ul>', '@') . '@',
-      $json[$contextual_links_id], $matches);
-
-    if ($found) {
-      $options = UrlHelper::parse($matches[1]);
-      return Url::fromUri('internal:/' . $options['path'], $options);
-    }
-
-    return FALSE;
-  }
-
-}
diff --git a/src/Tests/FlagTestBase.php b/src/Tests/FlagTestBase.php
index 2cdf91c..8433cf7 100644
--- a/src/Tests/FlagTestBase.php
+++ b/src/Tests/FlagTestBase.php
@@ -6,6 +6,8 @@ use Drupal\Core\StringTranslation\StringTranslationTrait;
 use Drupal\simpletest\WebTestBase;
 use Drupal\flag\Entity\Flag;
 use Drupal\Core\Template\Attribute;
+use Drupal\flag\Tests\FlagWithFormTrait;
+use Drupal\Tests\flag\Traits\FlagPermissionTrait;
 
 /**
  * Provides common methods for Flag tests.
@@ -14,6 +16,7 @@ abstract class FlagTestBase extends WebTestBase {
 
   use FlagCreateTrait;
   use FlagPermissionsTrait;
+  use FlagWithFormTrait;
   use StringTranslationTrait;
 
   /**
@@ -85,63 +88,6 @@ abstract class FlagTestBase extends WebTestBase {
   );
 
   /**
-   * Creates a flag entity using the admin UI.
-   *
-   * If you do not provide any bundles in $edit, all bundles for $entity_type
-   * are assumed.
-   *
-   * @param string|null $entity_type
-   *   (optional) A string containing the flaggable entity type, by default
-   *   'node'.
-   * @param array $edit
-   *   (optional) An array of form field names and values. If omitted, random
-   *   strings will be used for the flag ID, label, short and long text.
-   * @param string|null $link_type
-   *   (optional) A string containing the link type ID. Is omitted, assumes
-   *   'reload'.
-   *
-   * @return \Drupal\flag\FlagInterface|null
-   *   The created flag entity.
-   */
-  protected function createFlagWithForm($entity_type = 'node', $edit = [], $link_type = 'reload') {
-    // Submit the flag add page.
-    $this->drupalPostForm('admin/structure/flags/add', [
-      'flag_entity_type' => $this->getFlagType($entity_type),
-    ], $this->t('Continue'));
-
-    // Set the link type.
-    $this->drupalPostAjaxForm(NULL, ['link_type' => $link_type], 'link_type');
-
-    // Create an array of defaults.
-    $default_edit = [
-      'id' => strtolower($this->randomMachineName()),
-      'label' => $this->randomString(),
-      'flag_short' => $this->randomString(),
-      'unflag_short' => $this->randomString(),
-      'flag_long' => $this->randomString(16),
-      'unflag_long' => $this->randomString(16),
-      'flag_message' => $this->randomString(32),
-      'unflag_message' => $this->randomString(32),
-      'unflag_denied_text' => $this->randomString(),
-    ];
-
-    // Merge the default values with the edit array.
-    $final_edit = array_merge($default_edit, $edit);
-
-    // Submit the flag details form.
-    $this->drupalPostForm(NULL, $final_edit, $this->t('Create Flag'));
-
-    // Load the new flag we created.
-    $flag = Flag::load($final_edit['id']);
-
-    // Make sure that we actually did get a flag entity.
-    $this->assertTrue($flag instanceof Flag);
-
-    // Return the flag.
-    return $flag;
-  }
-
-  /**
    * Asserts that a contextual link placeholder with the given id exists.
    *
    * @see \Drupal\contextual\Tests\ContextualDynamicContextTest::assertContextualLinkPlaceHolder().
diff --git a/src/Tests/FlagWithFormTrait.php b/src/Tests/FlagWithFormTrait.php
new file mode 100644
index 0000000..e202fc5
--- /dev/null
+++ b/src/Tests/FlagWithFormTrait.php
@@ -0,0 +1,69 @@
+<?php
+
+namespace Drupal\flag\Tests;
+
+use Drupal\flag\Entity\Flag;
+
+/**
+ * Trait for creating Flags using web pages.
+ */
+trait FlagWithFormTrait {
+
+    /**
+   * Creates a flag entity using the admin UI.
+   *
+   * If you do not provide any bundles in $edit, all bundles for $entity_type
+   * are assumed.
+   *
+   * @param string|null $entity_type
+   *   (optional) A string containing the flaggable entity type, by default
+   *   'node'.
+   * @param array $edit
+   *   (optional) An array of form field names and values. If omitted, random
+   *   strings will be used for the flag ID, label, short and long text.
+   * @param string|null $link_type
+   *   (optional) A string containing the link type ID. Is omitted, assumes
+   *   'reload'.
+   *
+   * @return \Drupal\flag\FlagInterface|null
+   *   The created flag entity.
+   */
+  protected function createFlagWithForm($entity_type = 'node', $edit = [], $link_type = 'reload') {
+    // Submit the flag add page.
+    $this->drupalPostForm('admin/structure/flags/add', [
+      'flag_entity_type' => $this->getFlagType($entity_type),
+    ], 'Continue');
+
+    // Set the link type.
+    $this->drupalPostAjaxForm(NULL, ['link_type' => $link_type], 'link_type');
+
+    // Create an array of defaults.
+    $default_edit = [
+      'id' => strtolower($this->randomMachineName()),
+      'label' => $this->randomString(),
+      'flag_short' => $this->randomString(),
+      'unflag_short' => $this->randomString(),
+      'flag_long' => $this->randomString(16),
+      'unflag_long' => $this->randomString(16),
+      'flag_message' => $this->randomString(32),
+      'unflag_message' => $this->randomString(32),
+      'unflag_denied_text' => $this->randomString(),
+    ];
+
+    // Merge the default values with the edit array.
+    $final_edit = array_merge($default_edit, $edit);
+
+    // Submit the flag details form.
+    $this->drupalPostForm(NULL, $final_edit, 'Create Flag');
+
+    // Load the new flag we created.
+    $flag = Flag::load($final_edit['id']);
+
+    // Make sure that we actually did get a flag entity.
+    $this->assertTrue($flag instanceof Flag);
+
+    // Return the flag.
+    return $flag;
+  }
+
+}
\ No newline at end of file
diff --git a/tests/src/FunctionalJavascript/FlagContextualLinksTest.php b/tests/src/FunctionalJavascript/FlagContextualLinksTest.php
new file mode 100644
index 0000000..c798bf0
--- /dev/null
+++ b/tests/src/FunctionalJavascript/FlagContextualLinksTest.php
@@ -0,0 +1,184 @@
+<?php
+
+namespace Drupal\Tests\flag\FunctionalJavascript;
+
+use Drupal\flag\FlagInterface;
+use Drupal\Tests\flag\Traits\FlagWithFormTrait;
+use Drupal\Tests\flag\Traits\FlagPermissionsTrait;
+use Drupal\flag\Tests\FlagCreateTrait;
+use Drupal\Tests\flag\FunctionalJavascript\FlagTestBase;
+use Drupal\node\NodeInterface;
+use Behat\Mink\Element\ElementInterface;
+/**
+ * Test the contextual links with Reload link type.
+    *
+ * @group flag
+ */
+class FlagContextualLinksTest extends FlagTestBase {
+
+  use FlagWithFormTrait;
+  use FlagCreateTrait;
+  use FlagPermissionsTrait;
+
+  /**
+   * {@inheritdoc}
+   */
+  public static $modules = [
+    'contextual',
+    'field_ui',
+    'flag',
+    'node',
+    'system',
+    'user',
+  ];
+
+  /**
+   * The flag.
+   *
+   * @var \Drupal\flag\Flaginterface
+   */
+  protected $flag;
+
+  /**
+   * A user with Flag admin rights.
+   *
+   * @var AccountInterface
+   */
+  protected $adminUser;
+
+  /**
+   * An authenticated user to test flagging.
+   *
+   * @var \Drupal\user\Entity\User|false
+   */
+  protected $normalUser;
+
+  /**
+   * The node type to use in the test.
+   *
+   * @var string
+   */
+  protected $nodeType = 'article';
+
+  /**
+   * {@inheritdoc}
+   */
+  public function setUp() {
+    parent::setUp();
+
+    // Create content type.
+    $this->drupalCreateContentType(['type' => $this->nodeType]);
+
+    // Create the admin user.
+    $this->adminUser = $this->createUser([
+      'administer flags',
+      'administer flagging display',
+      'administer flagging fields',
+      'administer node display',
+      'administer modules',
+      'administer nodes',
+      'create ' . $this->nodeType . ' content',
+      'edit any ' . $this->nodeType . ' content',
+      'delete any ' . $this->nodeType . ' content',
+    ]);
+
+    // Create a regular user who will be flagging content.
+    $this->normalUser = $this->drupalCreateUser();
+
+    $this->drupalLogin($this->adminUser);
+
+    //  Create a flag with a Reload link type and enable contextual links display.
+    $edit = [
+      'id' => 'test_label_123',
+      'show_as_field' => FALSE,
+      'show_contextual_link' => TRUE,
+    ];
+    $this->flag = $this->createFlagWithForm('node', $edit);
+  }
+
+  /**
+   * Verify the behaviour and rendering of Flag links.
+   */
+  public function testFlagLinks() {
+    // Grant the flag permissions to the authenticated role.
+    $this->grantFlagPermissions($this->flag);
+
+    $node = $this->drupalCreateNode(['type' => $this->nodeType]);
+
+    // Login as normal user.
+    $this->drupalLogin($this->normalUser);
+
+    // Open node view page that renders Full view mode.
+    $this->drupalGet('node/' . $node->id());
+
+    // Expect the contextual links identifer for this node to contain flag_keys
+    // metadata related to the flag action.
+    $flag_contextual_links_id = 'node:node=' . $node->id() . ':changed=' . $node->getChangedTime() . '&flag_keys=' . $this->flag->id() . '-flag&langcode=en';
+    $page = $this->getSession()->getPage();
+    $flag_placeholder = $page->find('css', 'div[data-contextual-id="' . $flag_contextual_links_id . '"]');
+    $this->assertTrue($flag_placeholder->isValid(), format_string('Contextual link placeholder with id @id exists (flag).', ['@id' => $flag_contextual_links_id]));
+
+    // Open the contextual links pannel.
+    $flag_placeholder->find('css', 'button.trigger.focusable')->click();
+
+    // While the contextual links are being rendered check the structure is correct.
+    $this->assertPlaceholderMatches($flag_placeholder, $this->flag, $node);
+
+    // Click on the flag link embedded in the contextual links.
+    // Once the contextual list structure becomes visible.
+    $flag_placeholder->waitForElement('css', 'ul:visible');
+    $anchor = $flag_placeholder->findLink($this->flag->getMessage('flag'));
+    $anchor->click();
+
+    // Verify the contextual links are updated to contain unflag links.
+    $unflag_contextual_links_id = 'node:node=' . $node->id() . ':changed=' . $node->getChangedTime() . '&flag_keys=' . $this->flag->id() . '-unflag&langcode=en';
+    $unflag_placeholder = $this->assert_session()->waitForELement('css', 'div[data-contextual-id="' . $unflag_contextual_links_id . '"]');
+    assertTrue($unflag_placeholder->isValid(), format_string('Contextual link placeholder with id @id exists (unflag).', ['@id' => $unflag_contextual_links_id]));
+    $this->assertPlaceholderMatches($unflag_placeholder, $this->flag, $node);
+
+    // Login as normal user.
+    $this->drupalLogin($this->drupalCreateUser());
+
+    // Open node view page that renders Full view mode.
+    $this->drupalGet('node/' . $node->id());
+    $page2 = $this->getSession()->getPage();
+
+    // Expect the contextual links identifier for this node to contain flag_keys
+    // metadata related to the flag action.
+    $flag_placeholder2 = $page2->find('css', 'div[data-contextual-id="' . $flag_contextual_links_id . '"]');
+    $this->assertTrue($flag_placeholder2->isValid(), format_string('Contextual link placeholder with id @id exists.', ['@id' => $flag_contextual_links_id]));
+    $flag_action_url2 = $this->assertPlaceholderMatches($flag_placeholder2, $this->flag, $node);
+    $this->assertTrue($flag_action_url2, 'Flag link found in contextual links (normal user).');
+  }
+
+  /**
+   * Returns a flag action contextual link in the form of a URL.
+   *
+   * A contextual identifier is needed to fully specify the final markup.
+   *
+   * @param Behat\Mink\Element\ElementInterface $placeholder
+   *   The HTML elements wrapper holding the contextual0links
+   * @param \Drupal\flag\FlagInterface $flag
+   *   The flag associated with the contextual link.
+   * @param Drupal\node\NodeInterface $node
+   *   The node associated with the contextual link.
+   */
+  protected function assertPlaceholderMatches(ElementInterface $placeholder, FlagInterface $flag, NodeInterface $node) {
+
+    // @todo - tofix flag->id() is hardcoded as flag-test-label-123.
+    $anchor_element = $placeholder->find('css', 'ul.contextual-links li.flag-test-label-123 a');
+    $this->assertTrue($anchor_element->isValid(), 'Found expected contextual-links unsorted list structure.');
+
+    // Extract href from the page
+    $href = $anchor_element->getAttribute('href');
+
+    // Compute the expected link for a given flag and node combination.
+    $flag_link = $flag
+      ->getLinkTypePlugin()
+      ->getAsLink($flag, $node);
+
+    $flag_url = $flag_link->getUrl();
+    $this->assertEquals($href, $flag_url, 'A link was displayed macthing the given Flag and node pair.');
+  }
+
+}
diff --git a/tests/src/FunctionalJavascript/FlagTestBase.php b/tests/src/FunctionalJavascript/FlagTestBase.php
new file mode 100644
index 0000000..2c6649f
--- /dev/null
+++ b/tests/src/FunctionalJavascript/FlagTestBase.php
@@ -0,0 +1,51 @@
+<?php
+
+namespace Drupal\Tests\flag\FunctionalJavascript;
+
+use Drupal\flag\Tests\FlagCreateTrait;
+use Drupal\FunctionalJavascriptTests\JavascriptTestBase;
+use Drupal\Tests\flag\Traits\FlagWithFormTrait;
+use Drupal\Tests\flag\Traits\FlagPermissionsTrait;
+
+/**
+ * Base class for functional javascript tests related to the flag module.
+ *
+ * @group flag
+ */
+class FlagTestBase extends JavascriptTestBase {
+
+  use FlagWithFormTrait;
+  use FlagCreateTrait;
+  use FlagPermissionsTrait;
+
+  /**
+   * Asserts that a contextual link placeholder with the given id exists.
+   *
+   * @see \Drupal\contextual\Tests\ContextualDynamicContextTest::assertContextualLinkPlaceHolder().
+   *
+   * @param string $id
+   *   A contextual link id.
+   *
+   */
+  protected function assertContextualLinkPlaceHolder($id) {
+    $page = $this->assertSession()->getPage();
+    $found = $page->has('css', 'div[data-contextual-id="id"]');
+    assertTrue($found, format_string('Contextual link placeholder with id @id exists.', ['@id' => $id]));
+  }
+
+  /**
+   * Asserts that a contextual link placeholder with the given id exists.
+   *
+   * @see \Drupal\contextual\Tests\ContextualDynamicContextTest::assertContextualLinkPlaceHolder().
+   *
+   * @param string $id
+   *   A contextual link id.
+   *
+   */
+  protected function assertNoContextualLinkPlaceholder($id) {
+    $page = $this->assertSession()->getPage();
+    $found = $page->has('css', 'div[data-contextual-id="id"]');
+    assertFalse($found, format_string('Contextual link placeholder with id @id does not exists.', ['@id' => $id]));
+  }
+
+}
\ No newline at end of file
diff --git a/src/Tests/FlagPermissionsTrait.php b/tests/src/Traits/FlagPermissionsTrait.php
similarity index 97%
rename from src/Tests/FlagPermissionsTrait.php
rename to tests/src/Traits/FlagPermissionsTrait.php
index aec7853..a8206e7 100644
--- a/src/Tests/FlagPermissionsTrait.php
+++ b/tests/src/Traits/FlagPermissionsTrait.php
@@ -1,6 +1,6 @@
 <?php
 
-namespace Drupal\flag\Tests;
+namespace Drupal\Tests\flag\Traits;
 
 use Drupal\user\Entity\Role;
 use Drupal\user\RoleInterface;
diff --git a/tests/src/Traits/FlagWithFormTrait.php b/tests/src/Traits/FlagWithFormTrait.php
new file mode 100644
index 0000000..f67d884
--- /dev/null
+++ b/tests/src/Traits/FlagWithFormTrait.php
@@ -0,0 +1,70 @@
+<?php
+
+namespace Drupal\Tests\flag\Traits;
+
+use Drupal\flag\Entity\Flag;
+
+/**
+ * Trait for creating Flags using web pages.
+ */
+trait FlagWithFormTrait {
+
+    /**
+   * Creates a flag entity using the admin UI.
+   *
+   * If you do not provide any bundles in $edit, all bundles for $entity_type
+   * are assumed.
+   *
+   * @param string|null $entity_type
+   *   (optional) A string containing the flaggable entity type, by default
+   *   'node'.
+   * @param array $edit
+   *   (optional) An array of form field names and values. If omitted, random
+   *   strings will be used for the flag ID, label, short and long text.
+   * @param string|null $link_type
+   *   (optional) A string containing the link type ID. Is omitted, assumes
+   *   'reload'.
+   *
+   * @return \Drupal\flag\FlagInterface|null
+   *   The created flag entity.
+   */
+  protected function createFlagWithForm($entity_type = 'node', $edit = [], $link_type = 'reload') {
+    // Submit the flag add page.
+     // Verify default, non-modal behavior.
+    $this->drupalGet('admin/structure/flags/add');
+    $first_page = $this->getSession()->getPage();
+    // @todo Vary which radio button is pressed based on $entity_type.
+    $radio_input = $first_page->findField('flag_entity_type');
+    $radio_input->click();
+    $first_page->findButton('Continue')->click();
+
+    // Create an array of defaults.
+    $default_edit = [
+      'id' => strtolower($this->randomMachineName()),
+      'label' => $this->randomString(),
+      'flag_short' => $this->randomString(),
+      'unflag_short' => $this->randomString(),
+      'flag_long' => $this->randomString(16),
+      'unflag_long' => $this->randomString(16),
+      'flag_message' => $this->randomString(32),
+      'unflag_message' => $this->randomString(32),
+      'unflag_denied_text' => $this->randomString(),
+    ];
+
+    // Merge the default values with the edit array.
+    $final_edit = array_merge($default_edit, $edit);
+
+    // Submit the flag details form.
+    $this->drupalPostForm(NULL, $final_edit, 'Create Flag');
+
+    // Load the new flag we created.
+    $flag = Flag::load($final_edit['id']);
+
+    // Make sure that we actually did get a flag entity.
+    $this->assertTrue($flag instanceof Flag);
+
+    // Return the flag.
+    return $flag;
+  }
+
+}
