unki/lib/u_lib/src/models/jobs/misc.rs

use diesel_derive_enum::DbEnum;
use serde::{Deserialize, Serialize};
use strum::Display;

#[derive(Serialize, Deserialize, Clone, Debug)]
pub enum ManageAction {
    Ping,
    UpdateAvailable,
    JobsResultsRequest,
    Terminate,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)]
pub enum JobSchedule {
    Once,
    Permanent,
    //Scheduled
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum, Display)]
#[PgType = "JobState"]
#[DieselType = "Jobstate"]
pub enum JobState {
    Queued, // server created a job, but client didn't get it yet
    //Pending, // client got a job, but not running yet
    Running, // client is currently running a job
    Finished,
}

#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, DbEnum, Display)]
#[PgType = "JobType"]
#[DieselType = "Jobtype"]
pub enum JobType {
    Manage,
    Shell,
    Python,
}

#[derive(Clone, Debug)]
pub struct JobOutput {
    pub stdout: Vec<u8>,
    pub stderr: Vec<u8>,
}

impl JobOutput {
    const STREAM_BORDER: &'static str = "***";
    const STDOUT: &'static str = "STDOUT";
    const STDERR: &'static str = "STDERR";

    #[inline]
    fn create_delim(header: &'static str) -> Vec<u8> {
        format!(
            "<{border}{head}{border}>",
            border = Self::STREAM_BORDER,
            head = header
        )
        .into_bytes()
    }

    pub fn new() -> Self {
        Self {
            stdout: vec![],
            stderr: vec![],
        }
    }

    pub fn stdout(mut self, data: Vec<u8>) -> Self {
        self.stdout = data;
        self
    }

    pub fn stderr(mut self, data: Vec<u8>) -> Self {
        self.stderr = data;
        self
    }

    pub fn into_combined(self) -> Vec<u8> {
        let mut result: Vec<u8> = vec![];
        if self.stdout.len() > 0 {
            result.extend(Self::create_delim(Self::STDOUT));
            result.extend(self.stdout);
        }

        if self.stderr.len() > 0 {
            result.extend(Self::create_delim(Self::STDERR));
            result.extend(self.stderr);
        }
        result
    }

    pub fn from_combined(raw: &[u8]) -> Option<Self> {
        enum ParseFirst {
            Stdout,
            Stderr,
        }
        fn split_by_subslice<'s>(slice: &'s [u8], subslice: &[u8]) -> Option<(&'s [u8], &'s [u8])> {
            slice
                .windows(subslice.len())
                .position(|w| w == subslice)
                .map(|split_pos| {
                    let splitted = slice.split_at(split_pos);
                    (&splitted.0[..split_pos], &splitted.1[subslice.len()..])
                })
        }
        let splitter = |p: ParseFirst| {
            let (first_hdr, second_hdr) = match p {
                ParseFirst::Stdout => (Self::STDOUT, Self::STDERR),
                ParseFirst::Stderr => (Self::STDERR, Self::STDOUT),
            };
            let first_hdr = Self::create_delim(first_hdr);
            let second_hdr = Self::create_delim(second_hdr);
            split_by_subslice(raw, &first_hdr).map(|(_, p2)| {
                match split_by_subslice(p2, &second_hdr) {
                    Some((p2_1, p2_2)) => Self::new().stdout(p2_1.to_vec()).stderr(p2_2.to_vec()),
                    None => Self::new().stdout(p2.to_vec()),
                }
            })
        };
        splitter(ParseFirst::Stdout).or(splitter(ParseFirst::Stderr))
    }

    pub fn to_appropriate(&self) -> Vec<u8> {
        let mut result: Vec<u8> = vec![];
        if self.stdout.len() > 0 {
            result.extend(&self.stdout);
        }
        if self.stderr.len() > 0 {
            if result.len() > 0 {
                result.push(b'\n');
            }
            result.extend(&self.stderr);
        }
        if result.len() == 0 {
            result.extend(b"No data");
        }
        result
    }
}

#[cfg(test)]
mod tests {
    use crate::{models::JobOutput, utils::bytes_to_string};
    use test_case::test_case;

    const STDOUT: &str = "<***STDOUT***>";
    const STDERR: &str = "<***STDERR***>";

    #[test_case(
        "lol",
        "kek",
        &format!("{}lol{}kek", STDOUT, STDERR)
        ;"stdout stderr"
    )]
    #[test_case(
        "",
        "kek",
        &format!("{}kek", STDERR)
        ;"stderr"
    )]
    fn test_to_combined(stdout: &str, stderr: &str, result: &str) {
        let output = JobOutput::new()
            .stdout(stdout.as_bytes().to_vec())
            .stderr(stderr.as_bytes().to_vec());
        assert_eq!(&bytes_to_string(&output.into_combined()), result)
    }

    #[test_case(
        &format!("{}lal{}kik", STDOUT, STDERR),
        "lal\nkik"
        ;"stdout stderr"
    )]
    #[test_case(
        &format!("{}qeq", STDOUT),
        "qeq"
        ;"stdout"
    )]
    #[test_case(
        &format!("{}vev", STDERR),
        "vev"
        ;"stderr"
    )]
    fn test_from_combined(src: &str, result: &str) {
        let output = JobOutput::from_combined(src.as_bytes()).unwrap();
        assert_eq!(bytes_to_string(&output.to_appropriate()).trim(), result);
    }
}