--- a/testing/firefox-ui/harness/firefox_ui_harness/testcases.py
+++ b/testing/firefox-ui/harness/firefox_ui_harness/testcases.py
@@ -307,17 +307,17 @@ class UpdateTestCase(PuppeteerMixin, Mar
         finally:
             self.update_status['patch'] = self.patch_info
 
         if force_fallback:
             # Set the downloaded update into failed state
             self.software_update.force_fallback()
 
         # Restart Firefox to apply the downloaded update
-        self.restart()
+        self.restart(callback=lambda: about_window.deck.apply.button.click())
 
     def download_and_apply_forced_update(self):
         self.check_update_not_applied()
 
         # The update wizard dialog opens automatically after the restart but with a short delay
         dialog = Wait(self.marionette, ignored_exceptions=[NoSuchWindowException]).until(
             lambda _: self.puppeteer.windows.switch_to(lambda win: type(win) is UpdateWizardDialog)
         )
@@ -336,32 +336,36 @@ class UpdateTestCase(PuppeteerMixin, Mar
                 self.assertTrue(update_available,
                                 'Available update has been found')
 
                 # Download update and wait until it has been applied
                 self.download_update(about_window)
                 self.wait_for_update_applied(about_window)
 
             finally:
-                if about_window:
-                    self.update_status['patch'] = self.patch_info
+                self.update_status['patch'] = self.patch_info
 
+            # Restart Firefox to apply the downloaded fallback update
+            self.assertIsNotNone(about_window)
+            self.restart(callback=lambda: about_window.deck.apply.button.click())
+
+        # For a broken partial update, the software update window is used
         else:
             try:
                 self.assertEqual(dialog.wizard.selected_panel,
                                  dialog.wizard.error_patching)
 
                 # Start downloading the fallback update
                 self.download_update(dialog)
 
             finally:
                 self.update_status['patch'] = self.patch_info
 
-        # Restart Firefox to apply the update
-        self.restart()
+            # Restart Firefox to apply the downloaded fallback update
+            self.restart(callback=lambda: dialog.wizard.finish_button.click())
 
     def read_update_log(self):
         """Read the content of the update log file for the last update attempt."""
         path = os.path.join(os.path.dirname(self.software_update.staging_directory),
                             'last-update.log')
         try:
             with open(path, 'rb') as f:
                 return f.read().splitlines()