You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
162 lines
4.6 KiB
162 lines
4.6 KiB
3 years ago
|
use std::{time::SystemTime};
|
||
|
use u_lib::{
|
||
|
errors::UError,
|
||
|
models::{
|
||
|
jobs::{JobMeta},
|
||
|
ExecResult,
|
||
|
misc::JobType
|
||
|
},
|
||
|
builder::{JobBuilder, NamedJobBuilder}
|
||
|
};
|
||
|
|
||
|
type TestResult<R = ()> = Result<R, Box<dyn std::error::Error>>;
|
||
|
|
||
|
#[tokio::test]
|
||
|
async fn test_is_really_async() {
|
||
|
const SLEEP_SECS: u64 = 1;
|
||
|
let job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap();
|
||
|
let sleep_jobs: Vec<JobMeta> = (0..50).map(|_| job.clone()).collect();
|
||
|
let now = SystemTime::now();
|
||
|
JobBuilder::from_meta(sleep_jobs).unwrap_one().wait().await;
|
||
|
assert!(now.elapsed().unwrap().as_secs() < SLEEP_SECS + 2)
|
||
|
}
|
||
|
|
||
|
#[test_case(
|
||
|
"/bin/sh {}",
|
||
|
Some(b"echo test01 > /tmp/asd; cat /tmp/asd"),
|
||
|
"test01"
|
||
|
;"sh payload"
|
||
|
)]
|
||
|
#[test_case(
|
||
|
r#"/usr/bin/python -c 'print("test02")'"#,
|
||
|
None,
|
||
|
"test02"
|
||
|
;"python cmd"
|
||
|
)]
|
||
|
#[test_case(
|
||
|
"/{}",
|
||
|
Some(
|
||
|
br#"#!/bin/sh
|
||
|
TMPPATH=/tmp/lol
|
||
|
mkdir -p $TMPPATH
|
||
|
echo test03 > $TMPPATH/t
|
||
|
cat $TMPPATH/t"#
|
||
|
),
|
||
|
"test03"
|
||
|
;"sh multiline payload"
|
||
|
)]
|
||
|
#[test_case(
|
||
|
"/{} 'some msg as arg'",
|
||
|
Some(include_bytes!("../fixtures/echoer")),
|
||
|
"some msg as arg"
|
||
|
;"standalone binary with args"
|
||
|
)]
|
||
|
#[tokio::test]
|
||
|
async fn test_shell_job(cmd: &str, payload: Option<&[u8]>, expected_result: &str) -> TestResult {
|
||
|
let mut job = JobMeta::builder().with_shell(cmd);
|
||
|
if let Some(p) = payload {
|
||
|
job = job.with_payload(p);
|
||
|
}
|
||
|
let job = job.build().unwrap();
|
||
|
let job_result = JobBuilder::from_meta(job).unwrap_one().wait_one().await;
|
||
|
let result = unwrap_enum!(job_result, ExecResult::Assigned);
|
||
|
let result = result.to_string_result().unwrap();
|
||
|
assert_eq!(result.trim(), expected_result);
|
||
|
Ok(())
|
||
|
}
|
||
|
|
||
|
#[tokio::test]
|
||
|
async fn test_complex_load() -> TestResult {
|
||
|
const SLEEP_SECS: u64 = 1;
|
||
|
let now = SystemTime::now();
|
||
|
let longest_job = JobMeta::from_shell(format!("sleep {}", SLEEP_SECS)).unwrap();
|
||
|
let longest_job = JobBuilder::from_meta(longest_job).unwrap_one().spawn().await;
|
||
|
let ls = JobBuilder::from_meta(JobMeta::from_shell("ls")?).unwrap_one()
|
||
|
.wait_one()
|
||
|
.await;
|
||
|
let ls = unwrap_enum!(ls, ExecResult::Assigned);
|
||
|
assert_eq!(ls.retcode.unwrap(), 0);
|
||
|
let folders = ls.to_string_result().unwrap();
|
||
|
let subfolders_jobs: Vec<JobMeta> = folders
|
||
|
.lines()
|
||
|
.map(|f| JobMeta::from_shell(format!("ls {}", f)).unwrap())
|
||
|
.collect();
|
||
|
let ls_subfolders = JobBuilder::from_meta(subfolders_jobs)
|
||
|
.unwrap_one()
|
||
|
.wait()
|
||
|
.await;
|
||
|
for result in ls_subfolders {
|
||
|
let result = unwrap_enum!(result, ExecResult::Assigned);
|
||
|
assert_eq!(result.retcode.unwrap(), 0);
|
||
|
}
|
||
|
longest_job.wait().await;
|
||
|
assert_eq!(now.elapsed().unwrap().as_secs(), SLEEP_SECS);
|
||
|
Ok(())
|
||
|
}
|
||
|
/*
|
||
|
#[tokio::test]
|
||
|
async fn test_exec_multiple_jobs_nowait() -> UResult<()> {
|
||
|
const REPEATS: usize = 10;
|
||
|
let job = JobMeta::from_shell("whoami");
|
||
|
let sleep_jobs: Vec<JobMeta> = (0..=REPEATS).map(|_| job.clone()).collect();
|
||
|
build_jobs(sleep_jobs).spawn().await;
|
||
|
let mut completed = 0;
|
||
|
while completed < REPEATS {
|
||
|
let c = pop_completed().await.len();
|
||
|
if c > 0 {
|
||
|
completed += c;
|
||
|
println!("{}", c);
|
||
|
}
|
||
|
}
|
||
|
Ok(())
|
||
|
}
|
||
|
*/
|
||
|
#[tokio::test]
|
||
|
async fn test_failing_shell_job() -> TestResult {
|
||
|
let job = JobMeta::from_shell("lol_kek_puk")?;
|
||
|
let job_result = JobBuilder::from_meta(job)
|
||
|
.unwrap_one()
|
||
|
.wait_one()
|
||
|
.await;
|
||
|
let job_result = unwrap_enum!(job_result, ExecResult::Assigned);
|
||
|
let output = job_result.to_string_result().unwrap();
|
||
|
assert!(output.contains("No such file"));
|
||
|
assert!(job_result.retcode.is_none());
|
||
|
Ok(())
|
||
|
}
|
||
|
|
||
|
#[test_case(
|
||
|
"/bin/bash {}",
|
||
|
None,
|
||
|
"contains executable"
|
||
|
; "no binary"
|
||
|
)]
|
||
|
#[test_case(
|
||
|
"/bin/bash",
|
||
|
Some(b"whoami"),
|
||
|
"contains no executable"
|
||
|
; "no path to binary"
|
||
|
)]
|
||
|
#[tokio::test]
|
||
|
async fn test_job_building_failed(cmd: &str, payload: Option<&[u8]>, err_str: &str) -> TestResult {
|
||
|
let mut job = JobMeta::builder().with_shell(cmd);
|
||
|
if let Some(p) = payload {
|
||
|
job = job.with_payload(p);
|
||
|
}
|
||
|
let err = job.build().unwrap_err();
|
||
|
let err_msg = unwrap_enum!(err, UError::JobArgsError);
|
||
|
assert!(err_msg.contains(err_str));
|
||
|
Ok(())
|
||
|
}
|
||
|
|
||
|
#[tokio::test]
|
||
|
async fn test_different_job_types() -> TestResult {
|
||
|
let mut jobs = NamedJobBuilder::from_meta(vec![
|
||
|
("sleeper", JobMeta::from_shell("sleep 3")?),
|
||
|
("gatherer", JobMeta::builder().with_type(JobType::Manage).build()?)
|
||
|
]).wait().await;
|
||
|
let gathered = jobs.pop("gatherer");
|
||
|
assert_eq!(unwrap_enum!(gathered, ExecResult::Agent).alias, None);
|
||
|
Ok(())
|
||
|
}
|