fix:func to return correct error and fmt

Signed-off-by: Sayan Paul <[email protected]>

src/handler/mod.rs

@@ -1,19 +1,17 @@
 /// This module contains most of the low-level commands
 /// and grub variable modifications
 use anyhow::{bail, Result};
+
 use std::process::Command;
 use std::str;
 
 /// reboots the system if boot_counter is greater than 0 or can be forced too
 pub fn handle_reboot(force: bool) -> Result<()> {
     if !force {
-        match get_boot_counter() {
-            Some(t) => {
-                if t <= 0 {
-                    bail!("countdown ended, check greenboot-rollback status")
-                }
-            }
-            None => bail!("boot_counter is not set"),
+        if let Some(t) = get_boot_counter()? {
+            if t <= 0 {
+                bail!("countdown ended, check greenboot-rollback status")
+            };
         }
     }
     log::info!("restarting the system");
@@ -23,26 +21,29 @@ pub fn handle_reboot(force: bool) -> Result<()> {
 
 /// rollback to previous ostree deployment if boot counter is less than 0
 pub fn handle_rollback() -> Result<()> {
-    if let Some(t) = get_boot_counter() {
-        if t <= 0 {
+    if let Some(boot_counter) = get_boot_counter()? {
+        if boot_counter <= 0 {
             log::info!("Greenboot will now attempt to rollback");
             Command::new("rpm-ostree").arg("rollback").status()?;
             return Ok(());
-        } else {
-            bail!("Rollback not initiated, boot_counter {t}");
         }
     }
-    bail!("Rollback not initiated, boot_counter unset");
+    bail!("Rollback not initiated");
 }
 
 /// sets grub variable boot_counter if not set
 pub fn set_boot_counter(reboot_count: i32) -> Result<()> {
-    if let Some(current_counter) = get_boot_counter() {
-        log::info!("boot_counter={current_counter}");
-        return Ok(());
+    match get_boot_counter() {
+        Ok(Some(current_counter)) => {
+            log::info!("boot_counter={current_counter}");
+            bail!("counter already set");
+        }
+        Ok(None) => {
+            log::info!("setting boot counter");
+            set_grub_var("boot_counter", reboot_count)?;
+        }
+        Err(e) => bail!(e),
     }
-    log::info!("setting boot counter");
-    set_grub_var("boot_counter", reboot_count)?;
     Ok(())
 }
 
@@ -77,22 +78,13 @@ pub fn handle_motd(state: &str) -> Result<()> {
 }
 
 /// fetches boot_counter value, none if not set
-pub fn get_boot_counter() -> Option<i32> {
-    let grub_vars = Command::new("grub2-editenv").arg("-").arg("list").output();
-    if grub_vars.is_err() {
-        log::error!("{}", grub_vars.unwrap_err());
-        return None;
-    }
-    let grub_vars = grub_vars.unwrap();
-    let grub_vars = match str::from_utf8(&grub_vars.stdout[..]) {
-        Ok(vars) => vars.lines(),
-        Err(e) => {
-            log::error!("{e},unable to fetch grub variables");
-            return None;
-        }
-    };
-
-    for var in grub_vars {
+pub fn get_boot_counter() -> Result<Option<i32>> {
+    let grub_vars = Command::new("grub2-editenv")
+        .arg("-")
+        .arg("list")
+        .output()?;
+    let grub_vars = str::from_utf8(&grub_vars.stdout[..])?;
+    for var in grub_vars.lines() {
         let (k, v) = if let Some(kv) = var.split_once('=') {
             kv
         } else {
@@ -101,15 +93,11 @@ pub fn get_boot_counter() -> Option<i32> {
         if k != "boot_counter" {
             continue;
         }
-        match v.parse::<i32>() {
-            Ok(count) => return Some(count),
-            Err(_) => {
-                log::error!("boot_counter {v} , should be a valid integer");
-                return None;
-            }
-        }
+
+        let counter_value = v.parse::<i32>()?;
+        return Ok(Some(counter_value));
     }
-    None
+    Ok(None)
 }
 
 /// helper function to set any grub variable

src/main.rs

@@ -318,7 +318,7 @@ mod tests {
     fn test_boot_counter_set() {
         unset_boot_counter().ok();
         set_boot_counter(10).ok();
-        assert_eq!(get_boot_counter(), Some(10));
+        assert_eq!(get_boot_counter().unwrap(), Some(10));
         unset_boot_counter().ok();
     }
 
@@ -327,7 +327,7 @@ mod tests {
         unset_boot_counter().ok();
         set_boot_counter(10).ok();
         set_boot_counter(20).ok();
-        assert_eq!(get_boot_counter(), Some(10));
+        assert_eq!(get_boot_counter().unwrap(), Some(10));
         unset_boot_counter().ok();
     }
 
@@ -341,7 +341,7 @@ mod tests {
             .spawn()
             .context("Cannot create grub variable boot_counter");
         set_boot_counter(13).ok();
-        assert_eq!(get_boot_counter(), Some(13));
+        assert_eq!(get_boot_counter().unwrap(), Some(13));
         unset_boot_counter().ok();
     }