diff --git a/shell.nix b/shell.nix
index 949c095..571a14b 100644
--- a/shell.nix
+++ b/shell.nix
@@ -8,7 +8,7 @@ pkgs.mkShellNoCC {
       projectDir = ./.;
       preferWheels = true;
       editablePackageSources = {nix_background_upgrade = ./src;};
-      extraPackages = ps: with ps; [python-lsp-server];
+      extraPackages = ps: with ps; [mypy];
     })
     poetry
   ];
diff --git a/src/nix_background_upgrade/main.py b/src/nix_background_upgrade/main.py
index 5763cfd..334cc05 100644
--- a/src/nix_background_upgrade/main.py
+++ b/src/nix_background_upgrade/main.py
@@ -5,7 +5,14 @@ from typing import Optional
 
 from notifypy import Notify
 
-NIXOS_CONFIGURATION: str = getenv("NIXOS_CONFIGURATION", "")
+DEFAULT_NIXOS_CONFIGURATION: Path = Path("/etc/nixos")
+DEFAULT_HOME_MANAGER_CONFIGURATION: Path = Path("~/.config/home-manager/").resolve()
+NIXOS_CONFIGURATION: str = getenv(
+    "NIXOS_CONFIGURATION", str(DEFAULT_NIXOS_CONFIGURATION)
+)
+HOME_MANAGER_CONFIGURATION: str = getenv(
+    "HOME_MANAGER_CONFIGURATION", str(DEFAULT_HOME_MANAGER_CONFIGURATION)
+)
 
 
 def create_notification(source: Optional[str] = None) -> Notify:
@@ -20,45 +27,75 @@ def create_notification(source: Optional[str] = None) -> Notify:
     if source is None:
         source = "Nix process"
     return Notify(
-        default_notification_title="Message from {}".format(source),
+        default_notification_title="{}".format(source),
         default_notification_icon=Path(__file__).parent / "./nix-snowflake.png",
     )
 
 
+def _abstract_switch(
+    notification: Notify, command: str, flake: Optional[str] = None
+) -> int:
+    if flake is not None:
+        command += " --flake {}".format(flake)
+    try:
+        try:
+            process = subprocess.run(command.split())
+        except FileNotFoundError as e:
+            notification.message = e
+            return -1
+        try:
+            process.check_returncode()
+        except subprocess.CalledProcessError as e:
+            notification.message = e
+            return process.returncode
+        notification.message = "Success: {}".format(command)
+        return process.returncode
+    finally:
+        notification.send()
+
+
 def nixos_switch() -> int:
     """
     Switch NixOS system to new version
 
-    :returns: Error code of `sudo nixos-rebuild switch` call
+    :returns: Return code of `sudo nixos-rebuild switch` call
     """
-    notification = create_notification(source="nixos_switch")
-    try:
-        subprocess.run("which -a nixos-rebuild".split()).check_returncode()
-    except subprocess.CalledProcessError as e:
-        notification.message = e
-        notification.send()
-    return 0
+    return _abstract_switch(
+        notification=create_notification(source="nixos_switch"),
+        command="which -a nixos-rebuild",
+        flake=(
+            NIXOS_CONFIGURATION
+            if NIXOS_CONFIGURATION != str(DEFAULT_NIXOS_CONFIGURATION)
+            else None
+        ),
+    )
 
 
-def home_manager_switch() -> None:
+def home_manager_switch() -> int:
     """
     Switch home-manager to new version
 
-    :returns: None
+    :returns: Return code of `home-manager switch` call
     """
-    subprocess.Popen("which -a home-manager".split())
-    return
+    return _abstract_switch(
+        notification=create_notification(source="home_manager_switch"),
+        command="which -a home-manager",
+        flake=(
+            HOME_MANAGER_CONFIGURATION
+            if HOME_MANAGER_CONFIGURATION != str(DEFAULT_HOME_MANAGER_CONFIGURATION)
+            else None
+        ),
+    )
 
 
-def main(*args, **kwargs) -> None:
+def main() -> None:
     """
     Sends a test notification to check everything is working
 
     :returns: None
     """
-    notification = create_notification()
-    notification.message = "This is a test notification"
-    notification.send()
+    print(nixos_switch())
+    print(home_manager_switch())
     return