databuild/databuild/mock_job_run.rs

83 lines
2.1 KiB
Rust

use std::collections::HashMap;
pub struct MockJobRun {
sleep_ms: u64,
stdout_msg: String,
output_file: Option<OutputFile>,
exit_code: u8,
}
pub struct OutputFile {
path: String,
contents: String,
}
impl Default for MockJobRun {
fn default() -> Self {
Self {
sleep_ms: 0,
stdout_msg: "test executed".to_string(),
output_file: None,
exit_code: 0,
}
}
}
impl MockJobRun {
pub fn new() -> Self {
Self::default()
}
pub fn sleep_ms(mut self, val: u64) -> Self {
self.sleep_ms = val;
self
}
pub fn stdout_msg(mut self, val: &String) -> Self {
self.stdout_msg = val.into();
self
}
pub fn output_file(mut self, path: &String, contents: &String) -> Self {
self.output_file = Some(OutputFile {
path: path.to_string(),
contents: contents.to_string(),
});
self
}
pub fn exit_code(mut self, val: u8) -> Self {
self.exit_code = val;
self
}
pub fn to_env(&self) -> HashMap<String, String> {
let mut env = HashMap::new();
env.insert(
"DATABUILD_TEST_SLEEP_MS".to_string(),
self.sleep_ms.to_string(),
);
env.insert(
"DATABUILD_TEST_EXIT_CODE".to_string(),
self.exit_code.to_string(),
);
env.insert("DATABUILD_TEST_STDOUT".to_string(), self.stdout_msg.clone());
if let Some(output_file) = &self.output_file {
env.insert(
"DATABUILD_TEST_OUTPUT_FILE".to_string(),
output_file.path.clone(),
);
env.insert(
"DATABUILD_TEST_OUTPUT_CONTENTS".to_string(),
output_file.contents.clone(),
);
}
env
}
pub fn bin_path() -> String {
std::env::var("TEST_SRCDIR")
.map(|srcdir| format!("{}/_main/databuild/test/test_job_helper", srcdir))
.unwrap_or_else(|_| "bazel-bin/databuild/test/test_job_helper".to_string())
}
}