plazmoid
4 years ago
parent
5bd1760d10
commit
5734145e8f
36 changed files with 913 additions and 562 deletions
@ -1,3 +1,3 @@ |
||||
[build] |
||||
target = "x86_64-unknown-linux-gnu" # -musl" |
||||
|
||||
target = "x86_64-unknown-linux-musl" |
||||
rustflags = ["-C", "target-feature=+crt-static"] |
||||
|
||||
@ -1 +0,0 @@ |
||||
export DATABASE_URL=postgres://postgres:12348756@172.17.0.2/u_db |
||||
@ -0,0 +1,3 @@ |
||||
FROM centos:7 |
||||
|
||||
CMD yum update |
||||
@ -0,0 +1,88 @@ |
||||
// TODO:
|
||||
// поддержка питона
|
||||
// резолв адреса управляющего сервера через DoT
|
||||
// кроссплатформенность (реализовать интерфейс для винды и никсов)
|
||||
// проверка обнов
|
||||
// самоуничтожение
|
||||
|
||||
#[macro_use] |
||||
extern crate log; |
||||
extern crate env_logger; |
||||
|
||||
use std::env; |
||||
use tokio::time::{sleep, Duration}; |
||||
use u_lib::{ |
||||
api::ClientHandler, |
||||
builder::JobBuilder, |
||||
cache::JobCache, |
||||
executor::pop_completed, |
||||
models::{AssignedJob, ExecResult}, |
||||
UID, |
||||
//daemonize
|
||||
}; |
||||
|
||||
#[macro_export] |
||||
macro_rules! retry_until_ok { |
||||
( $body:expr ) => { |
||||
loop { |
||||
match $body { |
||||
Ok(r) => break r, |
||||
Err(e) => error!("{:?}", e), |
||||
}; |
||||
sleep(Duration::from_secs(5)).await; |
||||
} |
||||
}; |
||||
} |
||||
|
||||
pub async fn process_request(job_requests: Vec<AssignedJob>, client: &ClientHandler) { |
||||
if job_requests.len() > 0 { |
||||
for jr in &job_requests { |
||||
if !JobCache::contains(&jr.job_id) { |
||||
info!("Fetching job: {}", &jr.job_id); |
||||
let fetched_job = retry_until_ok!(client.get_jobs(Some(jr.job_id)).await) |
||||
.pop() |
||||
.unwrap(); |
||||
JobCache::insert(fetched_job); |
||||
} |
||||
} |
||||
info!( |
||||
"Scheduling jobs: \n{}", |
||||
job_requests |
||||
.iter() |
||||
.map(|j| j.job_id.to_string()) |
||||
.collect::<Vec<String>>() |
||||
.join("\n") |
||||
); |
||||
let mut builder = JobBuilder::from_request(job_requests); |
||||
let errors = builder.pop_errors(); |
||||
if errors.len() > 0 { |
||||
error!( |
||||
"Some errors encountered: \n{}", |
||||
errors |
||||
.iter() |
||||
.map(|j| j.to_string()) |
||||
.collect::<Vec<String>>() |
||||
.join("\n") |
||||
); |
||||
} |
||||
builder.unwrap_one().spawn().await; |
||||
} |
||||
} |
||||
|
||||
pub async fn run_forever() { |
||||
//daemonize();
|
||||
env_logger::init(); |
||||
let arg_ip = env::args().nth(1); |
||||
let instance = ClientHandler::new(arg_ip); |
||||
info!("Connecting to the server"); |
||||
loop { |
||||
let job_requests: Vec<AssignedJob> = |
||||
retry_until_ok!(instance.get_agent_jobs(Some(*UID)).await); |
||||
process_request(job_requests, &instance).await; |
||||
let result: Vec<ExecResult> = pop_completed().await.into_iter().collect(); |
||||
if result.len() > 0 { |
||||
retry_until_ok!(instance.report(&result).await) |
||||
} |
||||
sleep(Duration::from_secs(5)).await; |
||||
} |
||||
} |
||||
@ -1,89 +1,7 @@ |
||||
// TODO:
|
||||
// поддержка питона
|
||||
// резолв адреса управляющего сервера через DoT
|
||||
// кроссплатформенность (реализовать интерфейс для винды и никсов)
|
||||
// проверка обнов
|
||||
// самоуничтожение
|
||||
|
||||
#[macro_use] |
||||
extern crate log; |
||||
extern crate env_logger; |
||||
|
||||
use std::env; |
||||
use tokio::time::{sleep, Duration}; |
||||
use u_lib::{ |
||||
api::ClientHandler, |
||||
builder::JobBuilder, |
||||
cache::JobCache, |
||||
executor::pop_completed, |
||||
models::{AssignedJob, ExecResult}, |
||||
UID, |
||||
//daemonize
|
||||
}; |
||||
|
||||
#[macro_export] |
||||
macro_rules! retry_until_ok { |
||||
( $body:expr ) => { |
||||
loop { |
||||
match $body { |
||||
Ok(r) => break r, |
||||
Err(e) => error!("{:?}", e), |
||||
}; |
||||
sleep(Duration::from_secs(5)).await; |
||||
} |
||||
}; |
||||
} |
||||
|
||||
async fn process_request(job_requests: Vec<AssignedJob>, client: &ClientHandler) { |
||||
if job_requests.len() > 0 { |
||||
for jr in &job_requests { |
||||
if !JobCache::contains(&jr.job_id) { |
||||
info!("Fetching job: {}", &jr.job_id); |
||||
let fetched_job = retry_until_ok!(client.get_jobs(Some(jr.job_id)).await) |
||||
.pop() |
||||
.unwrap(); |
||||
JobCache::insert(fetched_job); |
||||
} |
||||
} |
||||
info!( |
||||
"Scheduling jobs: \n{}", |
||||
job_requests |
||||
.iter() |
||||
.map(|j| j.job_id.to_string()) |
||||
.collect::<Vec<String>>() |
||||
.join("\n") |
||||
); |
||||
let mut builder = JobBuilder::from_request(job_requests); |
||||
let errors = builder.pop_errors(); |
||||
if errors.len() > 0 { |
||||
error!( |
||||
"Some errors encountered: \n{}", |
||||
errors |
||||
.iter() |
||||
.map(|j| j.to_string()) |
||||
.collect::<Vec<String>>() |
||||
.join("\n") |
||||
); |
||||
} |
||||
builder.unwrap_one().spawn().await; |
||||
} |
||||
} |
||||
use tokio; |
||||
use u_agent::run_forever; |
||||
|
||||
#[tokio::main] |
||||
async fn main() { |
||||
//daemonize();
|
||||
env_logger::init(); |
||||
let arg_ip = env::args().nth(1); |
||||
let instance = ClientHandler::new(arg_ip); |
||||
info!("Connecting to the server"); |
||||
loop { |
||||
let job_requests: Vec<AssignedJob> = |
||||
retry_until_ok!(instance.get_agent_jobs(Some(*UID)).await); |
||||
process_request(job_requests, &instance).await; |
||||
let result: Vec<ExecResult> = pop_completed().await.into_iter().collect(); |
||||
if result.len() > 0 { |
||||
retry_until_ok!(instance.report(&result).await) |
||||
} |
||||
sleep(Duration::from_secs(5)).await; |
||||
} |
||||
run_forever().await; |
||||
} |
||||
|
||||
@ -0,0 +1,8 @@ |
||||
use u_agent::process_request; |
||||
|
||||
type TestResult<R = ()> = Result<R, Box<dyn std::error::Error>>; |
||||
|
||||
#[tokio::test] |
||||
async fn test_first_connection() -> TestResult { |
||||
Ok(()) |
||||
} |
||||
@ -0,0 +1,2 @@ |
||||
extern crate u_agent; |
||||
mod behaviour; |
||||
@ -0,0 +1,3 @@ |
||||
FROM centos:7 |
||||
|
||||
CMD yum update |
||||
@ -0,0 +1,153 @@ |
||||
mod db; |
||||
mod handlers; |
||||
|
||||
use warp::{body, Filter, Rejection, Reply}; |
||||
|
||||
#[macro_use] |
||||
extern crate log; |
||||
extern crate env_logger; |
||||
|
||||
#[macro_use] |
||||
extern crate mockall; |
||||
#[macro_use] |
||||
extern crate mockall_double; |
||||
|
||||
use db::UDB; |
||||
#[double] |
||||
use handlers::Endpoints; |
||||
use serde::de::DeserializeOwned; |
||||
use u_lib::{ |
||||
config::MASTER_PORT, |
||||
messaging::{AsMsg, BaseMessage}, |
||||
models::*, |
||||
}; |
||||
use uuid::Uuid; |
||||
|
||||
fn get_content<M>() -> impl Filter<Extract = (BaseMessage<'static, M>,), Error = Rejection> + Clone |
||||
where |
||||
M: AsMsg + Sync + Send + DeserializeOwned + 'static, |
||||
{ |
||||
body::content_length_limit(1024 * 64).and(body::json::<BaseMessage<M>>()) |
||||
} |
||||
|
||||
fn prefill_jobs() { |
||||
let agent_hello = JobMeta::builder() |
||||
.with_type(misc::JobType::Manage) |
||||
.with_alias("agent_hello") |
||||
.build() |
||||
.unwrap(); |
||||
UDB::lock_db().insert_jobs(&[agent_hello]).ok(); |
||||
} |
||||
|
||||
fn init() { |
||||
env_logger::init(); |
||||
prefill_jobs(); |
||||
} |
||||
|
||||
fn make_filters() -> impl Filter<Extract = (impl Reply,), Error = Rejection> + Clone { |
||||
let infallible_none = |_| async { Ok::<(Option<Uuid>,), std::convert::Infallible>((None,)) }; |
||||
|
||||
let get_agents = warp::get() |
||||
.and(warp::path("get_agents")) |
||||
.and( |
||||
warp::path::param::<Uuid>() |
||||
.map(Some) |
||||
.or_else(infallible_none), |
||||
) |
||||
.and_then(Endpoints::get_agents); |
||||
|
||||
let upload_jobs = warp::post() |
||||
.and(warp::path("upload_jobs")) |
||||
.and(get_content::<Vec<JobMeta>>()) |
||||
.and_then(Endpoints::upload_jobs); |
||||
|
||||
let get_jobs = warp::get() |
||||
.and(warp::path("get_jobs")) |
||||
.and( |
||||
warp::path::param::<Uuid>() |
||||
.map(Some) |
||||
.or_else(infallible_none), |
||||
) |
||||
.and_then(Endpoints::get_jobs); |
||||
|
||||
let get_agent_jobs = warp::get() |
||||
.and(warp::path("get_agent_jobs")) |
||||
.and( |
||||
warp::path::param::<Uuid>() |
||||
.map(Some) |
||||
.or_else(infallible_none), |
||||
) |
||||
.and_then(|uid| Endpoints::get_agent_jobs(uid, false)); |
||||
|
||||
let get_personal_jobs = warp::get() |
||||
.and(warp::path("get_agent_jobs")) |
||||
.and(warp::path::param::<Uuid>().map(Some)) |
||||
.and_then(|uid| Endpoints::get_agent_jobs(uid, true)); |
||||
|
||||
let del = warp::get() |
||||
.and(warp::path("del")) |
||||
.and(warp::path::param::<Uuid>()) |
||||
.and_then(Endpoints::del); |
||||
|
||||
let set_jobs = warp::post() |
||||
.and(warp::path("set_jobs")) |
||||
.and(warp::path::param::<Uuid>()) |
||||
.and(get_content::<Vec<Uuid>>()) |
||||
.and_then(Endpoints::set_jobs); |
||||
|
||||
let report = warp::post() |
||||
.and(warp::path("report")) |
||||
.and(get_content::<Vec<ExecResult>>().and_then(Endpoints::report)); |
||||
|
||||
let auth_token = warp::header::exact("authorization", "Bearer 123qwe"); |
||||
|
||||
let agent_zone = get_jobs.clone().or(get_personal_jobs).or(report); |
||||
|
||||
let auth_zone = auth_token.and( |
||||
get_agents |
||||
.or(get_jobs) |
||||
.or(upload_jobs) |
||||
.or(del) |
||||
.or(set_jobs) |
||||
.or(get_agent_jobs), |
||||
); |
||||
|
||||
agent_zone.or(auth_zone) |
||||
} |
||||
|
||||
pub async fn serve() { |
||||
init(); |
||||
let routes = make_filters(); |
||||
warp::serve(routes.with(warp::log("warp"))) |
||||
.run(([0, 0, 0, 0], MASTER_PORT)) |
||||
.await; |
||||
} |
||||
|
||||
#[cfg(test)] |
||||
mod tests { |
||||
use super::*; |
||||
use handlers::build_ok; |
||||
use mockall::predicate::*; |
||||
use test_case::test_case; |
||||
use warp::test::request; |
||||
|
||||
#[test_case(Some(Uuid::new_v4()))] |
||||
#[test_case(None => panics)] |
||||
#[tokio::test] |
||||
async fn test_get_agent_jobs_unauthorized(uid: Option<Uuid>) { |
||||
let mock = Endpoints::get_agent_jobs_context(); |
||||
mock.expect() |
||||
.with(eq(uid), eq(uid.is_some())) |
||||
.returning(|_, _| Ok(build_ok(""))); |
||||
request() |
||||
.path(&format!( |
||||
"/get_agent_jobs/{}", |
||||
uid.map(|u| u.simple().to_string()).unwrap_or(String::new()) |
||||
)) |
||||
.method("GET") |
||||
.filter(&make_filters()) |
||||
.await |
||||
.unwrap(); |
||||
mock.checkpoint() |
||||
} |
||||
} |
||||
@ -0,0 +1,2 @@ |
||||
DATABASE_URL=postgres://postgres:12348756@u_db/u_db |
||||
PG_PASSWORD=12348756 |
||||
@ -0,0 +1,67 @@ |
||||
version: "2.1" |
||||
|
||||
networks: |
||||
u_net: |
||||
ipam: |
||||
config: |
||||
- subnet: 10.16.0.0/24 |
||||
|
||||
services: |
||||
|
||||
u_server: |
||||
image: unki/u_server |
||||
networks: |
||||
u_net: |
||||
ipv4_address: 10.16.0.200 |
||||
volumes: |
||||
- ../target/x86_64-unknown-linux-musl/release/u_server:/u_server |
||||
- ./.env:/.env |
||||
command: /u_server |
||||
depends_on: |
||||
u_db: |
||||
condition: service_started |
||||
expose: |
||||
- '63714' |
||||
environment: |
||||
RUST_LOG: warp |
||||
|
||||
u_db: |
||||
image: postgres:13.3 |
||||
networks: |
||||
- u_net |
||||
expose: |
||||
- '5432' |
||||
environment: |
||||
- POSTGRES_PASSWORD=${PG_PASSWORD} |
||||
|
||||
u_agent_1: |
||||
image: unki/u_agent |
||||
networks: |
||||
- u_net |
||||
volumes: |
||||
- ../target/x86_64-unknown-linux-musl/release/u_agent:/u_agent |
||||
command: /u_agent 10.16.0.200 |
||||
depends_on: |
||||
u_server: |
||||
condition: service_started |
||||
|
||||
u_agent_2: |
||||
image: unki/u_agent |
||||
networks: |
||||
- u_net |
||||
volumes: |
||||
- ../target/x86_64-unknown-linux-musl/release/u_agent:/u_agent |
||||
command: /u_agent 10.16.0.200 |
||||
depends_on: |
||||
u_server: |
||||
condition: service_started |
||||
|
||||
tests_runner: |
||||
image: unki/tests_runner |
||||
networks: |
||||
- u_net |
||||
volumes: |
||||
- ../:/unki/ |
||||
depends_on: |
||||
u_server: |
||||
condition: service_started |
||||
@ -0,0 +1,76 @@ |
||||
import subprocess |
||||
from utils import * |
||||
|
||||
DOCKERFILES = { |
||||
'u_agent': { |
||||
'ctx': '../bin/u_agent', |
||||
}, |
||||
'u_server': { |
||||
'ctx': '../bin/u_server' |
||||
}, |
||||
'tests_runner': { |
||||
'ctx': './', |
||||
'dockerfile_prefix': 'tests_runner' |
||||
}, |
||||
} |
||||
|
||||
|
||||
def docker(args): |
||||
try: |
||||
cmd = ['docker'] + args |
||||
#log(f'Running docker cmd: {cmd}') |
||||
return subprocess.check_output( |
||||
cmd |
||||
) |
||||
except subprocess.CalledProcessError as e: |
||||
err(str(e)) |
||||
raise |
||||
|
||||
|
||||
def print_errors(errors): |
||||
err_msg = '\n'.join( |
||||
' {container}: {error}'.format(container=item['container'], |
||||
error=item['error']) |
||||
for item in errors) |
||||
|
||||
err('There are some errors in next containers:\n%s' % err_msg) |
||||
|
||||
|
||||
def check_state(containers): |
||||
errors = [] |
||||
for container in containers: |
||||
ret, out = subprocess.getstatusoutput( |
||||
'docker inspect --format \'{{ .State.Running }}\' %s' |
||||
% container) |
||||
out = out.strip() |
||||
if ret == 0: |
||||
if out == 'true': |
||||
continue |
||||
else: |
||||
errors.append({'container': container, |
||||
'error': 'Bad state: Running=%s' % out}) |
||||
else: |
||||
errors.append({'container': container, |
||||
'error': out}) |
||||
|
||||
return errors |
||||
|
||||
|
||||
def rebuild_images_if_needed(force_rebuild=False): |
||||
for img_name, data in DOCKERFILES.items(): |
||||
ctx = data['ctx'] |
||||
df_prefix = data.get('dockerfile_prefix') |
||||
df_suffix = 'Dockerfile' |
||||
img_name = f'unki/{img_name}' |
||||
log(f'Building docker image {img_name}') |
||||
cmd = [ |
||||
'build', |
||||
'-t', |
||||
img_name, |
||||
ctx, |
||||
] |
||||
if df_prefix: |
||||
cmd += ['-f', f'{ctx}/{df_prefix}.{df_suffix}'] |
||||
if force_rebuild: |
||||
cmd += ['--no-cache'] |
||||
docker(cmd) |
||||
@ -0,0 +1,61 @@ |
||||
import subprocess |
||||
from utils import * |
||||
from docker import docker, check_state, print_errors |
||||
|
||||
|
||||
class Compose: |
||||
ALL_CONTAINERS = [ |
||||
'u_agent_1', |
||||
'u_agent_2', |
||||
'u_server', |
||||
'u_db', |
||||
'tests_runner', |
||||
] |
||||
|
||||
def __init__(self): |
||||
self.cmd_container = 'tests_runner' |
||||
|
||||
def _call(self, *args): |
||||
subprocess.check_call([ |
||||
'docker-compose', |
||||
'--no-ansi', |
||||
] + list(args) |
||||
) |
||||
|
||||
def up(self): |
||||
log('Instanciating cluster') |
||||
self._call('up') |
||||
log('Ok') |
||||
|
||||
def down(self): |
||||
log('Shutting down cluster') |
||||
self._call('down') |
||||
log('Ok') |
||||
|
||||
def stop(self): |
||||
log('Stopping cluster') |
||||
self._call('stop') |
||||
log('Ok') |
||||
|
||||
def run(self, cmd): |
||||
container = self.cmd_container |
||||
log(f'Running command "{cmd}" in container {container}') |
||||
docker([ |
||||
'exec', |
||||
'-i', |
||||
f"'{container}'", |
||||
f"'{cmd}'" |
||||
]) |
||||
log('Ok') |
||||
|
||||
def is_alive(self): |
||||
log('Check if all containers are alive') |
||||
|
||||
errors = check_state(self.ALL_CONTAINERS) |
||||
log('Check done') |
||||
|
||||
if errors: |
||||
print_errors(errors) |
||||
raise TestsError('Error during `is_alive` check') |
||||
else: |
||||
log('All containers are alive') |
||||
@ -0,0 +1,33 @@ |
||||
import signal |
||||
import sys |
||||
from utils import * |
||||
from docker import rebuild_images_if_needed |
||||
from docker_compose import Compose |
||||
|
||||
|
||||
cluster = Compose() |
||||
|
||||
|
||||
def abort_handler(s, _): |
||||
warn(f'Received signal: {s}') |
||||
warn(f'Gracefully stopping...') |
||||
cluster.down() |
||||
|
||||
|
||||
def run_tests(): |
||||
for s in (signal.SIGTERM, signal.SIGINT, signal.SIGHUP): |
||||
signal.signal(s, abort_handler) |
||||
rebuild_images_if_needed() |
||||
try: |
||||
cluster.up() |
||||
cluster.is_alive() |
||||
cluster.run('cargo test --test integration') |
||||
except Exception as e: |
||||
err(e) |
||||
sys.exit(1) |
||||
finally: |
||||
cluster.down() |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
run_tests() |
||||
@ -0,0 +1,5 @@ |
||||
#!/bin/bash |
||||
set -e |
||||
CARGO=cargo |
||||
$CARGO build --release --target x86_64-unknown-linux-musl |
||||
python integration_tests.py |
||||
@ -0,0 +1,3 @@ |
||||
FROM rust:1.53 |
||||
|
||||
CMD rustup target add x86_64-unknown-linux-musl |
||||
@ -0,0 +1,32 @@ |
||||
from termcolor import colored |
||||
|
||||
__all__ = ['log', 'warn', 'err', 'TestsError'] |
||||
|
||||
|
||||
class TestsError(Exception): |
||||
pass |
||||
|
||||
|
||||
COLORS = { |
||||
'question': colored('[?]', 'yellow'), |
||||
'info': colored('[~]', 'green'), |
||||
'warning': colored('[!]', 'magenta'), |
||||
'error': colored('[X]', 'red'), |
||||
} |
||||
|
||||
|
||||
def warn(msg): |
||||
log(msg, log_lvl='w') |
||||
|
||||
|
||||
def err(msg): |
||||
log(msg, log_lvl='e') |
||||
|
||||
|
||||
def log(msg, log_lvl='i'): |
||||
for lvl, text in COLORS.items(): |
||||
if lvl.startswith(log_lvl): |
||||
print(f'{text} {msg}') |
||||
break |
||||
else: |
||||
ValueError('Unknown log level') |
||||
@ -0,0 +1,14 @@ |
||||
use std::path::PathBuf; |
||||
use std::process::Command; |
||||
|
||||
fn main() { |
||||
let echoer = PathBuf::from("./tests/fixtures/echoer"); |
||||
let mut echoer_src = echoer.clone(); |
||||
echoer_src.set_extension("rs"); |
||||
Command::new("rustc") |
||||
.args(&[echoer_src.to_str().unwrap(), "-o", echoer.to_str().unwrap()]) |
||||
.status() |
||||
.unwrap(); |
||||
println!("cargo:rerun-if-changed={}", echoer_src.display()); |
||||
println!("cargo:rerun-if-changed={}", echoer.display()); |
||||
} |
||||
@ -1,17 +0,0 @@ |
||||
/* |
||||
use std::fmt::Display; |
||||
use u_api_proc_macro::api_route; |
||||
use uuid::Uuid; |
||||
|
||||
struct Paths; |
||||
struct ClientHandler; |
||||
|
||||
#[test] |
||||
fn test_api_proc_macro() { |
||||
#[api_route("GET", Uuid)] |
||||
fn list<T: Display>(&self, msg: T) -> String {} |
||||
|
||||
#[api_route("POST", Uuid)] |
||||
fn report<T: Display>(&self, msg: T) -> String {} |
||||
} |
||||
*/ |
||||
Binary file not shown.
@ -1,161 +0,0 @@ |
||||
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(()) |
||||
} |
||||
@ -1,43 +0,0 @@ |
||||
use u_lib::{models::JobOutput, utils::vec_to_string}; |
||||
|
||||
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!(&vec_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!(vec_to_string(&output.to_appropriate()).trim(), result); |
||||
} |
||||
@ -1,10 +0,0 @@ |
||||
#[macro_use] |
||||
extern crate test_case; |
||||
|
||||
#[macro_use] |
||||
extern crate u_lib; |
||||
|
||||
mod jobs { |
||||
mod execution; |
||||
mod output; |
||||
} |
||||
@ -1,10 +0,0 @@ |
||||
#!/bin/bash |
||||
set -e |
||||
source $(dirname $0)/rootdir.sh #set ROOTDIR |
||||
docker run \ |
||||
-v $ROOTDIR:/volume \ |
||||
-v cargo-cache:/root/.cargo/registry \ |
||||
-w /volume \ |
||||
-it \ |
||||
clux/muslrust \ |
||||
cargo $@ |
||||
Loading…
Reference in new issue