tests: State restore logic

Cargo.lock

@@ -1285,7 +1285,9 @@ dependencies = [
 
 [[package]]
 name = "pueue-lib"
-version = "0.14.1-alpha.0"
+version = "0.14.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc01a3022618ca4d3681f67cdc2f5e897aec0af87413390e1bb69f36434665ce"
 dependencies = [
  "async-std",
  "async-tls",

tests/helper/daemon/control.rs

@@ -49,6 +49,15 @@ pub async fn pause_daemon(shared: &Shared) -> Result<Message> {
         .context("Failed to send Pause message")
 }
 
+/// Helper to pause the whole daemon
+pub async fn shutdown_daemon(shared: &Shared) -> Result<Message> {
+    let message = Message::DaemonShutdown(Shutdown::Graceful);
+
+    send_message(shared, message)
+        .await
+        .context("Failed to send Shutdown message")
+}
+
 pub async fn get_state(shared: &Shared) -> Result<Box<State>> {
     let response = send_message(shared, Message::Status).await?;
     match response {

tests/helper/daemon/helper.rs

@@ -0,0 +1,90 @@
+use std::fs::File;
+use std::io::Read;
+use std::path::Path;
+use std::process::Child;
+
+use anyhow::{anyhow, bail, Context, Result};
+use procfs::process::Process;
+use pueue_lib::network::message::Message;
+
+use super::sleep_ms;
+
+pub fn assert_success(message: Message) {
+    assert!(matches!(message, Message::Success(_)));
+}
+
+/// Get a daemon pid from a specific pueue directory.
+/// This function gives the daemon a little time to boot up, but ultimately crashes if it takes too
+/// long.
+pub fn get_pid(pueue_dir: &Path) -> Result<i32> {
+    let pid_file = pueue_dir.join("pueue.pid");
+
+    // Give the daemon about 1 sec to boot and create the pid file.
+    let tries = 10;
+    let mut current_try = 0;
+
+    while current_try < tries {
+        // The daemon didn't create the pid file yet. Wait for 100ms and try again.
+        if !pid_file.exists() {
+            sleep_ms(50);
+            current_try += 1;
+            continue;
+        }
+
+        let mut file = File::open(&pid_file).context("Couldn't open pid file")?;
+        let mut content = String::new();
+        file.read_to_string(&mut content)
+            .context("Couldn't write to file")?;
+
+        // The file has been created but not yet been written to.
+        if content.is_empty() {
+            sleep_ms(50);
+            current_try += 1;
+            continue;
+        }
+
+        let pid = content
+            .parse::<i32>()
+            .map_err(|_| anyhow!("Couldn't parse value: {}", content))?;
+        return Ok(pid);
+    }
+
+    bail!("Couldn't find pid file after about 1 sec.");
+}
+
+/// Waits for a daemon to shut down.
+/// This is done by waiting for the pid to disappear.
+pub fn wait_for_shutdown(pid: i32) -> Result<()> {
+    // Try to read the process. If this fails, the daemon already exited.
+    let process = match Process::new(pid) {
+        Ok(process) => process,
+        Err(_) => return Ok(()),
+    };
+
+    // Give the daemon about 1 sec to shutdown.
+    let tries = 10;
+    let mut current_try = 0;
+
+    while current_try < tries {
+        // Process is still alive, wait a little longer
+        if process.is_alive() {
+            sleep_ms(50);
+            current_try += 1;
+            continue;
+        }
+
+        return Ok(());
+    }
+
+    bail!("Couldn't find pid file after about 1 sec.");
+}
+
+pub fn kill_and_print_output(mut child: Child) -> Result<()> {
+    let _ = child.kill();
+    let output = child.wait_with_output()?;
+    println!("Stdout: \n{:?}", String::from_utf8_lossy(&output.stdout));
+
+    println!("Stderr: \n{:?}", String::from_utf8_lossy(&output.stderr));
+
+    Ok(())
+}

tests/helper/daemon/mod.rs

@@ -0,0 +1,9 @@
+mod control;
+mod helper;
+mod setup;
+
+pub use control::*;
+pub use helper::*;
+pub use setup::*;
+
+use super::*;

tests/helper/daemon/setup.rs

@@ -1,15 +1,16 @@
-use std::fs::File;
+use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
-use std::{collections::BTreeMap, io::Read};
+use std::process::{Child, Command, Stdio};
 
-use anyhow::{anyhow, bail, Context, Result};
+use anyhow::{bail, Context, Result};
+use assert_cmd::prelude::*;
 use tempdir::TempDir;
 use tokio::io::{self, AsyncWriteExt};
 
 use pueue_daemon_lib::run;
 use pueue_lib::settings::*;
 
-use super::sleep_ms;
+use super::{get_pid, sleep_ms};
 
 /// Spawn the daemon main logic in it's own async function.
 /// It'll be executed by the tokio multi-threaded executor.
@@ -36,8 +37,35 @@ pub fn boot_daemon(pueue_dir: &Path) -> Result<i32> {
     bail!("Daemon didn't boot after 1sec")
 }
 
-/// Internal helper function, which wraps the daemon main logic and prints any error.
-pub async fn run_and_handle_error(pueue_dir: PathBuf, test: bool) -> Result<()> {
+/// Spawn the daemon by calling the actual pueued binary.
+/// This function also checks for the pid file and the unix socket to pop-up.
+pub fn boot_standalone_daemon(pueue_dir: &Path) -> Result<Child> {
+    let child = Command::cargo_bin("pueued")?
+        .arg("--config")
+        .arg(pueue_dir.join("pueue.yml").to_str().unwrap())
+        .arg("-vvv")
+        .stdout(Stdio::piped())
+        .spawn()?;
+
+    let tries = 20;
+    let mut current_try = 0;
+
+    // Wait up to 1s for the unix socket to pop up.
+    let socket_path = pueue_dir.join("pueue.pid");
+    while current_try < tries {
+        sleep_ms(50);
+        if socket_path.exists() {
+            return Ok(child);
+        }
+
+        current_try += 1;
+    }
+
+    bail!("Daemon didn't boot in stand-alone mode after 1sec")
+}
+
+/// Internal helper function, which wraps the daemon main logic and prints any errors.
+async fn run_and_handle_error(pueue_dir: PathBuf, test: bool) -> Result<()> {
     if let Err(err) = run(Some(pueue_dir.join("pueue.yml")), test).await {
         let mut stdout = io::stdout();
         stdout
@@ -52,45 +80,6 @@ pub async fn run_and_handle_error(pueue_dir: PathBuf, test: bool) -> Result<()>
     Ok(())
 }
 
-/// Get a daemon pid from a specific pueue directory.
-/// This function gives the daemon a little time to boot up, but ultimately crashes if it takes too
-/// long.
-pub fn get_pid(pueue_dir: &Path) -> Result<i32> {
-    let pid_file = pueue_dir.join("pueue.pid");
-
-    // Give the daemon about 1 sec to boot and create the pid file.
-    let tries = 10;
-    let mut current_try = 0;
-
-    while current_try < tries {
-        // The daemon didn't create the pid file yet. Wait for 100ms and try again.
-        if !pid_file.exists() {
-            sleep_ms(50);
-            current_try += 1;
-            continue;
-        }
-
-        let mut file = File::open(&pid_file).context("Couldn't open pid file")?;
-        let mut content = String::new();
-        file.read_to_string(&mut content)
-            .context("Couldn't write to file")?;
-
-        // The file has been created but not yet been written to.
-        if content.is_empty() {
-            sleep_ms(50);
-            current_try += 1;
-            continue;
-        }
-
-        let pid = content
-            .parse::<i32>()
-            .map_err(|_| anyhow!("Couldn't parse value: {}", content))?;
-        return Ok(pid);
-    }
-
-    bail!("Couldn't find pid file after about 1 sec.");
-}
-
 pub fn base_setup() -> Result<(Settings, TempDir)> {
     // Create a temporary directory used for testing.
     let tempdir = TempDir::new("pueue_lib").unwrap();

tests/helper/mod.rs

@@ -3,13 +3,11 @@ use anyhow::Result;
 use tokio::io::{self, AsyncWriteExt};
 
 pub mod daemon;
-pub mod daemon_control;
 pub mod fixtures;
 pub mod network;
 pub mod wait;
 
 pub use daemon::*;
-pub use daemon_control::*;
 pub use network::*;
 pub use wait::*;
 

tests/unix/kill.rs

@@ -1,8 +1,11 @@
 use anyhow::Result;
-use pueue_lib::network::message::*;
-use pueue_lib::task::*;
+use pretty_assertions::assert_eq;
 use rstest::rstest;
 
+use pueue_lib::network::message::*;
+use pueue_lib::state::GroupStatus;
+use pueue_lib::task::*;
+
 use crate::helper::*;
 
 #[rstest]
@@ -13,7 +16,7 @@ use crate::helper::*;
         all: true,
         children: false,
         signal: None,
-    })
+    }), true
 )]
 #[case(
     Message::Kill(KillMessage {
@@ -22,7 +25,7 @@ use crate::helper::*;
         all: false,
         children: false,
         signal: None,
-    })
+    }), true
 )]
 #[case(
     Message::Kill(KillMessage {
@@ -31,7 +34,7 @@ use crate::helper::*;
         all: false,
         children: false,
         signal: None,
-    })
+    }), false
 )]
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 /// Test if killing running tasks works as intended.
@@ -40,7 +43,13 @@ use crate::helper::*;
 /// - Via the --all flag, which just kills everything.
 /// - Via the --group flag, which just kills everything in the default group.
 /// - Via specific ids.
-async fn test_kill_tasks(#[case] kill_message: Message) -> Result<()> {
+///
+/// If a whole group or everything is killed, the respective groups should also be paused!
+/// This is security measure to prevent unwanted task execution in an emergency.
+async fn test_kill_tasks(
+    #[case] kill_message: Message,
+    #[case] group_should_pause: bool,
+) -> Result<()> {
     let (settings, tempdir) = base_setup()?;
     let shared = &settings.shared;
     let _pid = boot_daemon(tempdir.path())?;
@@ -71,5 +80,10 @@ async fn test_kill_tasks(#[case] kill_message: Message) -> Result<()> {
         assert_eq!(task.result, Some(TaskResult::Killed));
     }
 
+    // Groups should be paused in specific modes.
+    if group_should_pause {
+        assert_eq!(state.groups.get("default").unwrap(), &GroupStatus::Paused);
+    }
+
     Ok(())
 }

tests/unix/mod.rs

@@ -1,4 +1,7 @@
 mod add;
 mod kill;
+/// Tests regarding state restoration from a previous run.
+mod restore;
+/// Tests for shutting down the daemon
 mod shutdown;
 mod start;

tests/unix/restore.rs

@@ -0,0 +1,59 @@
+use std::convert::TryInto;
+
+use anyhow::Result;
+use pretty_assertions::assert_eq;
+
+use pueue_lib::state::GroupStatus;
+
+use crate::helper::*;
+
+#[tokio::test]
+/// The daemon should start in the same state as before shutdown, if no tasks are queued.
+/// This function tests for the running state.
+async fn test_start_running() -> Result<()> {
+    let (settings, tempdir) = base_setup()?;
+    let shared = &settings.shared;
+
+    let child = boot_standalone_daemon(tempdir.path())?;
+
+    // Kill the daemon and wait for it to shut down.
+    assert_success(shutdown_daemon(&shared).await?);
+    wait_for_shutdown(child.id().try_into()?)?;
+
+    // Boot it up again
+    let mut child = boot_standalone_daemon(tempdir.path())?;
+
+    // Assert that the group is still running.
+    let state = get_state(shared).await?;
+    assert_eq!(state.groups.get("default").unwrap(), &GroupStatus::Running);
+
+    child.kill()?;
+    Ok(())
+}
+
+#[tokio::test]
+/// The daemon should start in the same state as before shutdown, if no tasks are queued.
+/// This function tests for the paused state.
+async fn test_start_paused() -> Result<()> {
+    let (settings, tempdir) = base_setup()?;
+    let shared = &settings.shared;
+
+    let child = boot_standalone_daemon(tempdir.path())?;
+
+    // Pause the daemon
+    pause_daemon(shared).await?;
+
+    // Kill the daemon and wait for it to shut down.
+    assert_success(shutdown_daemon(&shared).await?);
+    wait_for_shutdown(child.id().try_into()?)?;
+
+    // Boot it up again
+    let mut child = boot_standalone_daemon(tempdir.path())?;
+
+    // Assert that the group is still paused.
+    let state = get_state(shared).await?;
+    assert_eq!(state.groups.get("default").unwrap(), &GroupStatus::Paused);
+
+    child.kill()?;
+    Ok(())
+}

tests/unix/shutdown.rs

@@ -1,5 +1,4 @@
-use assert_cmd::prelude::*;
-use std::process::{Command, Stdio};
+use std::convert::TryInto;
 
 use anyhow::{Context, Result};
 
@@ -11,18 +10,11 @@ use crate::helper::*;
 fn test_ctrlc() -> Result<()> {
     let (_settings, tempdir) = base_setup()?;
 
-    let mut child = Command::cargo_bin("pueued")?
-        .arg("--config")
-        .arg(tempdir.path().join("pueue.yml").to_str().unwrap())
-        .arg("-vvv")
-        .stdout(Stdio::piped())
-        .spawn()?;
-
-    let pid = get_pid(tempdir.path())?;
+    let mut child = boot_standalone_daemon(tempdir.path())?;
 
     use nix::sys::signal::{kill, Signal};
     // Send SIGTERM signal to process via nix
-    let nix_pid = nix::unistd::Pid::from_raw(pid);
+    let nix_pid = nix::unistd::Pid::from_raw(child.id() as i32);
     kill(nix_pid, Signal::SIGTERM).context("Failed to send SIGTERM to pid")?;
 
     // Sleep for 500ms and give the daemon time to shut down
@@ -35,3 +27,23 @@ fn test_ctrlc() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+/// Spin up the daemon and send a graceful shutdown message afterwards.
+/// The daemon should shutdown normally and exit with a 0.
+async fn test_graceful_shutdown() -> Result<()> {
+    let (settings, tempdir) = base_setup()?;
+
+    let mut child = boot_standalone_daemon(tempdir.path())?;
+
+    // Kill the daemon gracefully and wait for it to shut down.
+    assert_success(shutdown_daemon(&settings.shared).await?);
+    wait_for_shutdown(child.id().try_into()?)?;
+
+    let result = child.try_wait();
+    assert!(matches!(result, Ok(Some(_))));
+    let code = result.unwrap().unwrap();
+    assert!(matches!(code.code(), Some(0)));
+
+    Ok(())
+}