From fd9ad17d7993a4a65a39f6c481dc5012ea3504e6 Mon Sep 17 00:00:00 2001
From: June <june@jsts.xyz>
Date: Sun, 9 Jun 2024 00:26:04 +0200
Subject: [PATCH] feature: add --reboot flag for remote operations (except for
 "reboot")

Also make the "reboot" operation basically an alias for "boot --reboot"
then and simplify the code of the operations.run function, since it
doesn't need to convert between operation and actual_operation anymore.
---
 src/infra_rebuild/cli/__init__.py        | 19 ++++++++++++-------
 src/infra_rebuild/operations/__init__.py | 15 ++++-----------
 2 files changed, 16 insertions(+), 18 deletions(-)

diff --git a/src/infra_rebuild/cli/__init__.py b/src/infra_rebuild/cli/__init__.py
index 38030ea..f8b5f6c 100644
--- a/src/infra_rebuild/cli/__init__.py
+++ b/src/infra_rebuild/cli/__init__.py
@@ -2,6 +2,8 @@ import click
 
 from infra_rebuild import operations
 
+reboot_option_help_text = "Reboot the target hosts after running the operation."
+
 
 @click.group(context_settings={"help_option_names": ["-h", "--help"]}, invoke_without_command=True)
 @click.version_option(prog_name="infra-rebuild")
@@ -28,24 +30,27 @@ def build_vm_with_bootloader(hosts):
 
 
 @infra_rebuild.command()
+@click.option("--reboot", is_flag=True, help=reboot_option_help_text)
 @click.argument("hosts", nargs=-1)
-def switch(hosts):
-    operations.run("switch", hosts)
+def switch(hosts, reboot):
+    operations.run("switch", hosts, reboot)
 
 
 @infra_rebuild.command()
+@click.option("--reboot", is_flag=True, help=reboot_option_help_text)
 @click.argument("hosts", nargs=-1)
-def boot(hosts):
-    operations.run("boot", hosts)
+def boot(hosts, reboot):
+    operations.run("boot", hosts, reboot)
 
 
 @infra_rebuild.command()
+@click.option("--reboot", is_flag=True, help=reboot_option_help_text)
 @click.argument("hosts", nargs=-1)
-def test(hosts):
-    operations.run("test", hosts)
+def test(hosts, reboot):
+    operations.run("test", hosts, reboot)
 
 
 @infra_rebuild.command()
 @click.argument("hosts", nargs=-1)
 def reboot(hosts):
-    operations.run("reboot", hosts)
+    operations.run("boot", hosts, True)
diff --git a/src/infra_rebuild/operations/__init__.py b/src/infra_rebuild/operations/__init__.py
index 0ae726e..b8d3f28 100644
--- a/src/infra_rebuild/operations/__init__.py
+++ b/src/infra_rebuild/operations/__init__.py
@@ -6,19 +6,12 @@ import sys
 from infra_rebuild import msg
 
 
-def run(operation, hosts):
-    act_remotely = False
-    reboot = False
+def run(operation, hosts, reboot=False):  # noqa: FBT002 - having reboot as a Boolean positional argument is fine I think # fmt: skip
     match operation:
         case "build" | "build-vm" | "build-vm-with-bootloader":
-            actual_operation = operation
+            act_remotely = False
         case "switch" | "boot" | "test":
             act_remotely = True
-            actual_operation = operation
-        case "reboot":
-            act_remotely = True
-            actual_operation = "boot"
-            reboot = True
         case _:
             msg.error("Internal Error: The given operation isn't valid.")
             sys.exit(1)
@@ -29,9 +22,9 @@ def run(operation, hosts):
             continue
 
         if act_remotely:
-            remote(actual_operation, host, reboot)
+            remote(operation, host, reboot)
         else:
-            local(actual_operation, host)
+            local(operation, host)
 
 
 def local(operation, host):