diff --git a/modules/video_embed_media/tests/src/Functional/UpgradePathTest.php b/modules/video_embed_media/tests/src/Functional/UpgradePathTest.php
index 9ac3b2e..ff9fbfd 100644
--- a/modules/video_embed_media/tests/src/Functional/UpgradePathTest.php
+++ b/modules/video_embed_media/tests/src/Functional/UpgradePathTest.php
@@ -34,7 +34,10 @@ class UpgradePathTest extends BrowserTestBase {
   /**
    * Test the upgrade path.
    */
-  public function testMediaBundleCreation() {
+  public function testModuleUpgradePath() {
+
+    $this->rebuildContainer();
+
     $this->drupalLogin($this->createAdminUser());
 
     // Create a media_entity_embeddable_video bundle and field.
@@ -61,12 +64,16 @@ class UpgradePathTest extends BrowserTestBase {
       'name[0][value]' => 'Test Media Entity',
     ], t('Save'));
 
+    $this->rebuildContainer();
+
     // Install video_embed_field.
     $this->drupalGet('admin/modules');
     $this->submitForm([
       'modules[Video Embed Field][video_embed_media][enable]' => '1',
     ], t('Install'));
 
+    $this->rebuildContainer();
+
     $this->assertUpgradeComplete();
 
     // Uninstall the module and ensure everything is still okay.
@@ -83,12 +90,12 @@ class UpgradePathTest extends BrowserTestBase {
    * Assert the upgrade was successful.
    */
   protected function assertUpgradeComplete() {
-    // Ensure the new type is selected.
-    $this->drupalGet('admin/structure/media/manage/embeddable_bundle');
-    $this->assertTrue(!empty($this->getSession()->getPage()->find('xpath', '//option[@value="video_embed_field" and @selected="selected"]')), 'The media type was updated.');
     // Ensure the media entity has updated values.
     $this->drupalGet('media/1/edit');
     $this->assertEquals($this->getSession()->getPage()->find('css', 'input[name="field_video_text_field[0][value]"]')->getValue(), 'https://www.youtube.com/watch?v=gnERPdAiuSo', 'Field values were copied.');
+    // Ensure the new type is selected.
+    $this->drupalGet('admin/structure/media/manage/embeddable_bundle');
+    $this->assertSession()->fieldValueEquals('type', 'video_embed_field');
   }
 
 }