use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum LogLevel {
Trace,
Debug,
Info,
Warn,
Error,
}
impl LogLevel {
#[must_use]
pub fn as_str(&self) -> &'static str {
match self {
Self::Trace => "TRACE",
Self::Debug => "DEBUG",
Self::Info => "INFO",
Self::Warn => "WARN",
Self::Error => "ERROR",
}
}
}
pub struct MockHost {
config: BTreeMap<String, String>,
kv: BTreeMap<(String, String), Vec<u8>>,
secrets: BTreeMap<String, String>,
logs: Vec<(LogLevel, String)>,
http_responses: BTreeMap<String, MockHttpResponse>,
}
#[derive(Debug, Clone)]
pub struct MockHttpResponse {
pub status: u16,
pub headers: BTreeMap<String, String>,
pub body: Vec<u8>,
}
impl MockHttpResponse {
#[must_use]
pub fn new(status: u16) -> Self {
Self {
status,
headers: BTreeMap::new(),
body: Vec::new(),
}
}
#[must_use]
pub fn ok(body: impl AsRef<[u8]>) -> Self {
Self {
status: 200,
headers: BTreeMap::new(),
body: body.as_ref().to_vec(),
}
}
#[must_use]
pub fn json(body: impl AsRef<[u8]>) -> Self {
let mut headers = BTreeMap::new();
headers.insert("content-type".to_string(), "application/json".to_string());
Self {
status: 200,
headers,
body: body.as_ref().to_vec(),
}
}
#[must_use]
pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.insert(key.into(), value.into());
self
}
#[must_use]
pub fn with_body(mut self, body: impl AsRef<[u8]>) -> Self {
self.body = body.as_ref().to_vec();
self
}
#[must_use]
pub fn with_status(mut self, status: u16) -> Self {
self.status = status;
self
}
}
impl MockHost {
#[must_use]
pub fn new() -> Self {
Self {
config: BTreeMap::new(),
kv: BTreeMap::new(),
secrets: BTreeMap::new(),
logs: Vec::new(),
http_responses: BTreeMap::new(),
}
}
pub fn set_config(&mut self, key: &str, value: &str) -> &mut Self {
self.config.insert(key.to_string(), value.to_string());
self
}
pub fn set_configs(&mut self, configs: &[(&str, &str)]) -> &mut Self {
for (key, value) in configs {
self.config.insert((*key).to_string(), (*value).to_string());
}
self
}
#[must_use]
pub fn get_config(&self, key: &str) -> Option<&String> {
self.config.get(key)
}
pub fn set_kv(&mut self, bucket: &str, key: &str, value: &[u8]) -> &mut Self {
self.kv
.insert((bucket.to_string(), key.to_string()), value.to_vec());
self
}
pub fn set_kv_string(&mut self, bucket: &str, key: &str, value: &str) -> &mut Self {
self.set_kv(bucket, key, value.as_bytes())
}
#[must_use]
pub fn get_kv(&self, bucket: &str, key: &str) -> Option<&Vec<u8>> {
self.kv.get(&(bucket.to_string(), key.to_string()))
}
#[must_use]
pub fn get_kv_string(&self, bucket: &str, key: &str) -> Option<String> {
self.get_kv(bucket, key)
.and_then(|v| core::str::from_utf8(v).ok())
.map(ToString::to_string)
}
pub fn delete_kv(&mut self, bucket: &str, key: &str) -> &mut Self {
self.kv.remove(&(bucket.to_string(), key.to_string()));
self
}
#[must_use]
pub fn list_kv_keys(&self, bucket: &str) -> Vec<String> {
self.kv
.keys()
.filter(|(b, _)| b == bucket)
.map(|(_, k)| k.clone())
.collect()
}
pub fn set_secret(&mut self, name: &str, value: &str) -> &mut Self {
self.secrets.insert(name.to_string(), value.to_string());
self
}
pub fn set_secrets(&mut self, secrets: &[(&str, &str)]) -> &mut Self {
for (name, value) in secrets {
self.secrets
.insert((*name).to_string(), (*value).to_string());
}
self
}
#[must_use]
pub fn get_secret(&self, name: &str) -> Option<&String> {
self.secrets.get(name)
}
pub fn set_http_response(&mut self, url: &str, response: MockHttpResponse) -> &mut Self {
self.http_responses.insert(url.to_string(), response);
self
}
#[must_use]
pub fn get_http_response(&self, url: &str) -> Option<&MockHttpResponse> {
self.http_responses.get(url)
}
pub fn log(&mut self, level: LogLevel, message: &str) -> &mut Self {
self.logs.push((level, message.to_string()));
self
}
#[must_use]
pub fn logs(&self) -> &[(LogLevel, String)] {
&self.logs
}
#[must_use]
pub fn logs_at_level(&self, level: LogLevel) -> Vec<&String> {
self.logs
.iter()
.filter(|(l, _)| *l == level)
.map(|(_, msg)| msg)
.collect()
}
pub fn clear_logs(&mut self) {
self.logs.clear();
}
#[must_use]
pub fn has_log(&self, level: LogLevel, contains: &str) -> bool {
self.logs
.iter()
.any(|(l, msg)| *l == level && msg.contains(contains))
}
#[must_use]
pub fn has_log_containing(&self, contains: &str) -> bool {
self.logs.iter().any(|(_, msg)| msg.contains(contains))
}
#[must_use]
pub fn log_count(&self, level: LogLevel) -> usize {
self.logs.iter().filter(|(l, _)| *l == level).count()
}
#[must_use]
pub fn total_log_count(&self) -> usize {
self.logs.len()
}
pub fn assert_log(&self, level: LogLevel, contains: &str) {
assert!(
self.has_log(level, contains),
"Expected log at level {:?} containing '{}', but found: {:?}",
level,
contains,
self.logs
);
}
pub fn assert_no_log(&self, level: LogLevel, contains: &str) {
assert!(
!self.has_log(level, contains),
"Expected no log at level {level:?} containing '{contains}', but found one",
);
}
pub fn assert_no_errors(&self) {
let errors: Vec<_> = self.logs_at_level(LogLevel::Error);
assert!(
errors.is_empty(),
"Expected no error logs, but found: {errors:?}",
);
}
pub fn reset(&mut self) {
self.config.clear();
self.kv.clear();
self.secrets.clear();
self.logs.clear();
self.http_responses.clear();
}
}
impl Default for MockHost {
fn default() -> Self {
Self::new()
}
}
#[macro_export]
macro_rules! zlayer_test {
($name:ident, $setup:expr, $test:expr) => {
#[test]
fn $name() {
let mut host = $crate::testing::MockHost::new();
$setup(&mut host);
$test(&host);
}
};
}
#[macro_export]
macro_rules! zlayer_async_test {
($name:ident, $setup:expr, $test:expr) => {
#[tokio::test]
async fn $name() {
let mut host = $crate::testing::MockHost::new();
$setup(&mut host);
$test(&host).await;
}
};
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_mock_host_config() {
let mut host = MockHost::new();
host.set_config("key1", "value1")
.set_config("key2", "value2");
assert_eq!(host.get_config("key1"), Some(&"value1".to_string()));
assert_eq!(host.get_config("key2"), Some(&"value2".to_string()));
assert_eq!(host.get_config("key3"), None);
}
#[test]
fn test_mock_host_configs_batch() {
let mut host = MockHost::new();
host.set_configs(&[("a", "1"), ("b", "2"), ("c", "3")]);
assert_eq!(host.get_config("a"), Some(&"1".to_string()));
assert_eq!(host.get_config("b"), Some(&"2".to_string()));
assert_eq!(host.get_config("c"), Some(&"3".to_string()));
}
#[test]
fn test_mock_host_kv() {
let mut host = MockHost::new();
host.set_kv("bucket1", "key1", &[1, 2, 3])
.set_kv_string("bucket1", "key2", "hello");
assert_eq!(host.get_kv("bucket1", "key1"), Some(&vec![1, 2, 3]));
assert_eq!(
host.get_kv_string("bucket1", "key2"),
Some("hello".to_string())
);
assert_eq!(host.get_kv("bucket1", "key3"), None);
}
#[test]
fn test_mock_host_kv_list_keys() {
let mut host = MockHost::new();
host.set_kv_string("cache", "user:1", "alice")
.set_kv_string("cache", "user:2", "bob")
.set_kv_string("other", "data", "value");
let keys = host.list_kv_keys("cache");
assert_eq!(keys.len(), 2);
assert!(keys.contains(&"user:1".to_string()));
assert!(keys.contains(&"user:2".to_string()));
}
#[test]
fn test_mock_host_kv_delete() {
let mut host = MockHost::new();
host.set_kv_string("bucket", "key", "value");
assert!(host.get_kv("bucket", "key").is_some());
host.delete_kv("bucket", "key");
assert!(host.get_kv("bucket", "key").is_none());
}
#[test]
fn test_mock_host_secrets() {
let mut host = MockHost::new();
host.set_secret("api_key", "secret123")
.set_secrets(&[("db_pass", "pass456"), ("token", "tok789")]);
assert_eq!(host.get_secret("api_key"), Some(&"secret123".to_string()));
assert_eq!(host.get_secret("db_pass"), Some(&"pass456".to_string()));
assert_eq!(host.get_secret("token"), Some(&"tok789".to_string()));
assert_eq!(host.get_secret("missing"), None);
}
#[test]
fn test_mock_host_logging() {
let mut host = MockHost::new();
host.log(LogLevel::Info, "Application started");
host.log(LogLevel::Warn, "Connection slow");
host.log(LogLevel::Error, "Database connection failed");
assert_eq!(host.logs().len(), 3);
assert!(host.has_log(LogLevel::Info, "started"));
assert!(host.has_log(LogLevel::Warn, "slow"));
assert!(host.has_log(LogLevel::Error, "failed"));
assert!(!host.has_log(LogLevel::Debug, "anything"));
}
#[test]
fn test_mock_host_log_filtering() {
let mut host = MockHost::new();
host.log(LogLevel::Info, "info1");
host.log(LogLevel::Info, "info2");
host.log(LogLevel::Error, "error1");
assert_eq!(host.log_count(LogLevel::Info), 2);
assert_eq!(host.log_count(LogLevel::Error), 1);
assert_eq!(host.log_count(LogLevel::Debug), 0);
assert_eq!(host.total_log_count(), 3);
let info_logs = host.logs_at_level(LogLevel::Info);
assert_eq!(info_logs.len(), 2);
}
#[test]
fn test_mock_host_clear_logs() {
let mut host = MockHost::new();
host.log(LogLevel::Info, "message");
assert_eq!(host.logs().len(), 1);
host.clear_logs();
assert_eq!(host.logs().len(), 0);
}
#[test]
fn test_mock_host_has_log_containing() {
let mut host = MockHost::new();
host.log(LogLevel::Info, "Processing request ID=123");
assert!(host.has_log_containing("ID=123"));
assert!(host.has_log_containing("Processing"));
assert!(!host.has_log_containing("Error"));
}
#[test]
fn test_mock_host_assert_log() {
let mut host = MockHost::new();
host.log(LogLevel::Info, "Operation completed successfully");
host.assert_log(LogLevel::Info, "completed");
host.assert_no_log(LogLevel::Error, "failed");
}
#[test]
fn test_mock_host_assert_no_errors() {
let mut host = MockHost::new();
host.log(LogLevel::Info, "All good");
host.log(LogLevel::Warn, "Minor issue");
host.assert_no_errors();
}
#[test]
#[should_panic(expected = "Expected no error logs")]
fn test_mock_host_assert_no_errors_fails() {
let mut host = MockHost::new();
host.log(LogLevel::Error, "Something went wrong");
host.assert_no_errors();
}
#[test]
fn test_mock_host_http_response() {
let mut host = MockHost::new();
host.set_http_response(
"https://api.example.com/data",
MockHttpResponse::json(r#"{"result": "ok"}"#),
);
let response = host
.get_http_response("https://api.example.com/data")
.unwrap();
assert_eq!(response.status, 200);
assert_eq!(
response.headers.get("content-type"),
Some(&"application/json".to_string())
);
}
#[test]
fn test_mock_http_response_builder() {
let response = MockHttpResponse::new(404)
.with_header("x-custom", "value")
.with_body(b"Not Found");
assert_eq!(response.status, 404);
assert_eq!(response.headers.get("x-custom"), Some(&"value".to_string()));
assert_eq!(response.body, b"Not Found");
}
#[test]
fn test_mock_host_reset() {
let mut host = MockHost::new();
host.set_config("key", "value")
.set_kv_string("bucket", "key", "data")
.set_secret("secret", "shhh")
.log(LogLevel::Info, "message")
.set_http_response("https://example.com", MockHttpResponse::ok(b""));
host.reset();
assert!(host.get_config("key").is_none());
assert!(host.get_kv("bucket", "key").is_none());
assert!(host.get_secret("secret").is_none());
assert!(host.logs().is_empty());
assert!(host.get_http_response("https://example.com").is_none());
}
#[test]
fn test_log_level_as_str() {
assert_eq!(LogLevel::Trace.as_str(), "TRACE");
assert_eq!(LogLevel::Debug.as_str(), "DEBUG");
assert_eq!(LogLevel::Info.as_str(), "INFO");
assert_eq!(LogLevel::Warn.as_str(), "WARN");
assert_eq!(LogLevel::Error.as_str(), "ERROR");
}
#[test]
fn test_mock_host_default() {
let host = MockHost::default();
assert!(host.logs().is_empty());
}
}