From 0c32d0b3a853e7a06643895c8acecd7a6521f413 Mon Sep 17 00:00:00 2001
From: Dave Curylo <dave@curylo.org>
Date: Fri, 7 Jun 2024 11:54:59 +0000
Subject: [PATCH] Adds support for 'power_state'

---
 FsCloudInit/Builders.fs                       | 24 ++++++++++++++++
 FsCloudInit/CloudConfig.fs                    | 28 +++++++++++++++++++
 FsCloudInitTests/BuilderTests.fs              | 12 ++++++++
 FsCloudInitTests/TestContent/power-state.yaml |  4 +++
 README.md                                     | 13 +++++++++
 5 files changed, 81 insertions(+)
 create mode 100644 FsCloudInitTests/TestContent/power-state.yaml

diff --git a/FsCloudInit/Builders.fs b/FsCloudInit/Builders.fs
index 3f2a522..5b419b5 100644
--- a/FsCloudInit/Builders.fs
+++ b/FsCloudInit/Builders.fs
@@ -102,6 +102,26 @@ module Builders =
 
     let aptSource = AptSourceBuilder()
 
+    type PowerStateBuilder() =
+        member _.Yield _ = PowerState.Default
+
+        [<CustomOperation "delay">]
+        member _.DelayStateChange(powerState, delay) = { powerState with Delay = delay }
+
+        [<CustomOperation "mode">]
+        member _.Mode(powerState, mode) = { powerState with Mode = mode }
+
+        [<CustomOperation "message">]
+        member _.Message(powerState, message) = { powerState with Message = message }
+
+        [<CustomOperation "timeout">]
+        member _.Timeout(powerState, timeout) = { powerState with Timeout = timeout }
+
+        [<CustomOperation "condition">]
+        member _.Condition(powerState, condition) = { powerState with Condition = condition }
+
+    let powerState = PowerStateBuilder()
+
     type UbuntuAdvantageBuilder() =
         member _.Yield _ = UbuntuAdvantage.Default
 
@@ -245,6 +265,10 @@ module Builders =
             { cloudConfig with
                 FinalMessage = Some message }
 
+        [<CustomOperation "power_state">]
+        member _.PowerState(cloudConfig: CloudConfig, powerState: PowerState) =
+            { cloudConfig with PowerState = Some powerState }
+
         [<CustomOperation "run_commands">]
         member _.RunCommands(cloudConfig: CloudConfig, commands: string seq seq) =
             let cmdList = commands |> Seq.map List.ofSeq |> List.ofSeq
diff --git a/FsCloudInit/CloudConfig.fs b/FsCloudInit/CloudConfig.fs
index 212fc41..31dc7c2 100644
--- a/FsCloudInit/CloudConfig.fs
+++ b/FsCloudInit/CloudConfig.fs
@@ -101,6 +101,31 @@ type RunCmd =
         match this with
         | RunCmd commands -> commands |> Seq.map Seq.ofList
 
+type PowerState =
+    { Delay: int
+      Mode: string
+      Message: string
+      Timeout: Nullable<int>
+      Condition: string }
+    
+    static member Default =
+        { Delay = 0
+          Mode = null
+          Message = null
+          Timeout = Nullable()
+          Condition = null }
+
+module PowerState =
+
+    module Mode =
+        [<Literal>]
+        let PowerOff = "powerooff"
+        [<Literal>]
+        let Reboot = "reboot"
+        [<Literal>]
+        let Halt = "halt"
+
+
 type UbuntuAdvantage =
     { Token: string
       Enable: string seq }
@@ -167,6 +192,7 @@ type CloudConfig =
       PackageUpdate: bool option
       PackageUpgrade: bool option
       PackageRebootIfRequired: bool option
+      PowerState: PowerState option
       RunCmd: RunCmd option
       UbuntuAdvantage: UbuntuAdvantage option
       Users: User seq
@@ -179,6 +205,7 @@ type CloudConfig =
           PackageUpdate = None
           PackageUpgrade = None
           PackageRebootIfRequired = None
+          PowerState = None
           RunCmd = None
           UbuntuAdvantage = None
           Users = []
@@ -190,6 +217,7 @@ type CloudConfig =
            Packages = this.Packages |> Seq.map (fun p -> p.Model) |> Serialization.serializableSeq
            PackageUpdate = this.PackageUpdate |> Option.toNullable
            PackageUpgrade = this.PackageUpgrade |> Option.toNullable
+           PowerState = this.PowerState |> Option.defaultValue Unchecked.defaultof<PowerState>
            Runcmd = this.RunCmd |> Option.map (fun runCmd -> runCmd.Model) |> Option.toObj
            UbuntuAdvantage = this.UbuntuAdvantage |> Option.defaultValue Unchecked.defaultof<UbuntuAdvantage>
            Users =
diff --git a/FsCloudInitTests/BuilderTests.fs b/FsCloudInitTests/BuilderTests.fs
index 070e7f5..1ed0433 100644
--- a/FsCloudInitTests/BuilderTests.fs
+++ b/FsCloudInitTests/BuilderTests.fs
@@ -116,6 +116,18 @@ let tests =
               |> Writer.write
               |> matchExpectedAt "final-message.yaml"
           }
+          test "Set power state to reboot" {
+              cloudConfig { 
+                  power_state (
+                      powerState {
+                          mode PowerState.Mode.Reboot
+                          message "Done with installation. Rebooting now."
+                      }
+                  )
+              }
+              |> Writer.write
+              |> matchExpectedAt "power-state.yaml"
+          }
           test "Run commands with cloudConfig builder" {
               cloudConfig {
                   run_commands
diff --git a/FsCloudInitTests/TestContent/power-state.yaml b/FsCloudInitTests/TestContent/power-state.yaml
new file mode 100644
index 0000000..4e376ba
--- /dev/null
+++ b/FsCloudInitTests/TestContent/power-state.yaml
@@ -0,0 +1,4 @@
+#cloud-config
+power_state:
+  mode: reboot
+  message: Done with installation. Rebooting now.
diff --git a/README.md b/README.md
index b19feb7..fe806f2 100644
--- a/README.md
+++ b/README.md
@@ -213,6 +213,19 @@ cloudConfig {
 |> Writer.write
 ```
 
+#### Reboot after initialized
+
+```f#
+cloudConfig { 
+    power_state (
+        powerState {
+            mode PowerState.Mode.Reboot
+            message "Done with installation. Rebooting now."
+        }
+    )
+}
+```
+
 #### Attach to Ubuntu Pro Subscription
 
 ```f#