Update package and restart service before anything (#22)

pi_top_usb_setup/pages/run_setup.py

@@ -17,7 +17,11 @@
     NotEnoughSpaceException,
 )
 from pi_top_usb_setup.system_updater import SystemUpdater
-from pi_top_usb_setup.utils import get_package_version, restart_service_and_skip_updates
+from pi_top_usb_setup.utils import (
+    RestartingSystemdService,
+    get_package_version,
+    restart_service_and_skip_user_confirmation_dialog,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -121,6 +125,9 @@ def run_setup(self):
             self._complete_onboarding()
 
             self.state.update({"run_state": RunStates.DONE})
+        except RestartingSystemdService:
+            logger.warning("Restarting systemd service, exiting ...")
+            return
         except Exception as e:
             logger.error(f"{e}")
 
@@ -170,12 +177,37 @@ def _update_system(self):
                 ),
                 on_error=lambda message: logger.error(f"{message}"),
             )
+
             version_before_update = get_package_version("pi-top-usb-setup")
+            logger.info(
+                f"Before update, 'pi-top-usb-setup' version is {version_before_update}"
+            )
+
             logger.info("Starting system update")
             self.state.update({"run_state": RunStates.UPDATING_SYSTEM})
+
+            # Update sources
             updater.update()
+
+            # Upgrade pi-top-usb-setup package first
+            updater.upgrade_package("pi-top-usb-setup")
+
+            # Restart service if it was updated
+            version_after_update = get_package_version("pi-top-usb-setup")
+            if version_before_update != version_after_update:
+                logger.warning(
+                    f"Package 'pi-top-usb-setup' was updated from '{version_before_update}' to '{version_after_update}', restarting app..."
+                )
+                restart_service_and_skip_user_confirmation_dialog(
+                    mount_point=self.fs.SETUP_FOLDER
+                )
+                raise RestartingSystemdService
+
+            # Upgrade system
             updater.upgrade()
             logger.info("Finished updating")
+        except RestartingSystemdService:
+            raise
         except NotAnAptRepository as e:
             logger.warning(f"{e}")
             return
@@ -185,11 +217,6 @@ def _update_system(self):
             )
             raise Exception(f"Update Error: {e}")
 
-        # Restart service if it was updated
-        if version_before_update != get_package_version("pi-top-usb-setup"):
-            logger.info("Package 'pi-top-usb-setup' was updated, restarting app...")
-            restart_service_and_skip_updates()
-
     def _configure_device(self):
         self.state.update({"run_state": RunStates.CONFIGURING_DEVICE})
         try:

pi_top_usb_setup/system_updater.py

@@ -85,5 +85,8 @@ def updates_env():
     def update(self) -> None:
         self._run_cmd("sudo apt-get update")
 
+    def upgrade_package(self, package_name) -> None:
+        self._run_cmd(f"sudo apt-get install -y {package_name}")
+
     def upgrade(self) -> None:
         self._run_cmd("sudo apt-get dist-upgrade -y")

pi_top_usb_setup/utils.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import shutil
+import signal
 import tarfile
 from pathlib import Path
 from queue import Queue
@@ -15,6 +16,10 @@
 logger = logging.getLogger(__name__)
 
 
+class RestartingSystemdService(Exception):
+    pass
+
+
 def systemctl(command: str, name: str, timeout=10) -> Optional[str]:
     try:
         cmd = f"systemctl {command} {name}.service"
@@ -173,15 +178,15 @@ def terminate(process):
         return exit_code
 
 
-def restart_service_and_skip_updates():
+def restart_service_and_skip_user_confirmation_dialog(mount_point: str):
     # Start an instance with arguments; these should be encoded
     encoded_args = run_command(
-        "systemd-escape -- '--skip-dialog --skip-update'", timeout=5
+        f"systemd-escape -- '{mount_point} --skip-dialog'", timeout=5
     )
     systemctl("start", f"pt-usb-setup@'{encoded_args}'")
 
     # Stop this instance
-    systemctl("stop", "pt-usb-setup")
+    os.kill(os.getpid(), signal.SIGINT)
 
 
 def get_linux_distro():