--- a/testing/firefox-ui/harness/firefox_ui_harness/testcases.py
+++ b/testing/firefox-ui/harness/firefox_ui_harness/testcases.py
@@ -300,17 +300,17 @@ class UpdateTestCase(PuppeteerMixin, Mar
         finally:
             self.updates[self.current_update_index]['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(callback=lambda: about_window.deck.apply.button.click())
+        self.restart()
 
     def download_and_apply_forced_update(self):
         # 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)
         )
 
         # In case of a broken complete update the about window has to be used
@@ -327,34 +327,33 @@ 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:
-                self.updates[self.current_update_index]['patch'] = self.patch_info
-
-            # Restart Firefox to apply the forced update
-            self.restart(callback=lambda: about_window.deck.apply.button.click())
+                if about_window:
+                    self.updates[self.current_update_index]['patch'] = self.patch_info
 
         else:
             try:
                 self.assertEqual(dialog.wizard.selected_panel,
                                  dialog.wizard.error_patching)
 
                 # Start downloading the fallback update
                 self.download_update(dialog)
+                dialog.close()
 
             finally:
                 self.updates[self.current_update_index]['patch'] = self.patch_info
 
-            # Restart Firefox to apply the forced update
-            self.restart(callback=lambda: dialog.wizard.finish_button.click())
+        # Restart Firefox to apply the update
+        self.restart()
 
     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()

--- a/testing/puppeteer/firefox/firefox_puppeteer/mixins.py
+++ b/testing/puppeteer/firefox/firefox_puppeteer/mixins.py
@@ -60,17 +60,20 @@ class PuppeteerMixin(object):
             self.browser.tabbar.close_all_tabs([self.browser.tabbar.tabs[0]])
             self.browser.tabbar.tabs[0].switch_to()
 
     def restart(self, **kwargs):
         """Restart Firefox and re-initialize data.
 
         :param flags: Specific restart flags for Firefox
         """
-        self.marionette.restart(in_app=not kwargs.get('clean'), **kwargs)
+        if kwargs.get('clean'):
+            self.marionette.restart(clean=True)
+        else:
+            self.marionette.restart(in_app=True)
 
         # Ensure that we always have a valid browser instance available
         self.browser = self.puppeteer.windows.switch_to(lambda win: type(win) is BrowserWindow)
 
     def setUp(self, *args, **kwargs):
         super(PuppeteerMixin, self).setUp(*args, **kwargs)
 
         self._start_handle_count = len(self.marionette.window_handles)