diff --git a/src/Tests/EntityEmbedUpdateHookTest.php b/src/Tests/EntityEmbedUpdateHookTest.php
index 15dc6c5..eff365d 100644
--- a/src/Tests/EntityEmbedUpdateHookTest.php
+++ b/src/Tests/EntityEmbedUpdateHookTest.php
@@ -24,6 +24,20 @@ class EntityEmbedUpdateHookTest extends UpdatePathTestBase {
   /**
    * {@inheritdoc}
    */
+  protected function setUp() {
+    parent::setUp();
+    $entity_browser_type = $this->container
+      ->get('entity_type.manager')
+      ->getDefinition('embed_button');
+
+    $this->container
+      ->get('entity.last_installed_schema.repository')
+      ->setLastInstalledDefinition($entity_browser_type);
+  }
+
+  /**
+   * {@inheritdoc}
+   */
   protected function doSelectionTest() {
     parent::doSelectionTest();
     $this->assertRaw('8002 -   Updates the default mode settings.');