use super::types::{
DebugResponsePreview, IsResponse, ResponseMode, RiftResponseExtension, RiftScriptConfig,
StubResponse,
};
use crate::behaviors::{apply_decorate, HasRepeatBehavior, RequestContext};
use crate::imposter::Predicate;
use std::collections::HashMap;
fn truncate_with_ellipsis(text: &str, max_len: usize) -> String {
if text.len() <= max_len {
return text.to_string();
}
let end = text.floor_char_boundary(max_len);
format!("{}...", &text[..end])
}
impl HasRepeatBehavior for StubResponse {
fn get_repeat(&self) -> Option<u32> {
match self {
StubResponse::Is { behaviors, .. } => behaviors
.as_ref()
.and_then(|b| b.get("repeat"))
.and_then(|r| r.as_u64())
.map(|r| r as u32),
StubResponse::RiftScript { .. } => None,
_ => None,
}
}
}
pub fn create_response_preview(response: &StubResponse) -> DebugResponsePreview {
match response {
StubResponse::Is { is, .. } => {
let body_preview = is.body.as_ref().map(|b| match b {
serde_json::Value::String(s) => truncate_with_ellipsis(s, 500),
other => {
let json = serde_json::to_string(other).unwrap_or_default();
truncate_with_ellipsis(&json, 500)
}
});
let headers = if is.headers.is_empty() {
None
} else {
Some(
is.headers
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect(),
)
};
DebugResponsePreview {
response_type: "is".to_string(),
status_code: Some(is.status_code),
headers,
body_preview,
}
}
StubResponse::Proxy { proxy, .. } => DebugResponsePreview {
response_type: "proxy".to_string(),
status_code: None,
headers: None,
body_preview: Some(format!("Proxy to: {}", proxy.to)),
},
StubResponse::Inject { inject, .. } => DebugResponsePreview {
response_type: "inject".to_string(),
status_code: None,
headers: None,
body_preview: Some(format!(
"JavaScript inject: {}",
truncate_with_ellipsis(inject, 50)
)),
},
StubResponse::Fault { fault, .. } => DebugResponsePreview {
response_type: "fault".to_string(),
status_code: None,
headers: None,
body_preview: Some(format!("Fault: {fault}")),
},
StubResponse::RiftScript { rift } => {
let script_info = if rift.script.is_some() {
"Rift script response"
} else if rift.fault.is_some() {
"Rift fault injection"
} else {
"Rift extension response"
};
DebugResponsePreview {
response_type: "_rift".to_string(),
status_code: None,
headers: None,
body_preview: Some(script_info.to_string()),
}
}
}
}
#[allow(clippy::type_complexity)]
pub fn execute_stub_response_with_rift(
response: &StubResponse,
) -> Option<(
u16,
HashMap<String, String>,
String,
Option<serde_json::Value>,
Option<RiftResponseExtension>,
ResponseMode,
bool,
)> {
match response {
StubResponse::Is {
is,
behaviors,
rift,
} => {
let mut headers = is.headers.clone();
let mode = is.mode.clone();
let body = is
.body
.as_ref()
.map(|b| {
if b.is_string() {
b.as_str().unwrap_or("").to_string()
} else {
if !headers.contains_key("content-type")
&& !headers.contains_key("Content-Type")
{
headers
.insert("Content-Type".to_string(), "application/json".to_string());
}
serde_json::to_string(b).unwrap_or_default()
}
})
.unwrap_or_default();
Some((
is.status_code,
headers,
body,
behaviors.clone(),
rift.clone(),
mode,
false,
))
}
StubResponse::Fault { fault } => Some((
0,
HashMap::new(),
fault.clone(),
None,
None,
ResponseMode::Text,
true,
)),
StubResponse::Proxy { .. } => None,
StubResponse::Inject { .. } => None,
StubResponse::RiftScript { .. } => None,
}
}
pub fn get_rift_script_config(response: &StubResponse) -> Option<RiftScriptConfig> {
match response {
StubResponse::RiftScript { rift } => rift.script.clone(),
_ => None,
}
}
pub fn create_stub_from_proxy_response(
predicates: Vec<serde_json::Value>,
status: u16,
headers: &[(String, String)],
body: &[u8],
latency_ms: Option<u64>,
decorate_fn: Option<String>,
recorded_from: Option<String>,
) -> super::types::Stub {
let response_headers: HashMap<String, String> = {
let merged = crate::util::merge_headers_to_map(headers);
merged
.into_iter()
.filter(|(k, _)| !crate::util::is_hop_by_hop_header(k))
.collect()
};
let (body_value, is_binary) = crate::util::encode_body_for_stub(body);
let mode = if is_binary {
ResponseMode::Binary
} else {
ResponseMode::Text
};
let is_response = IsResponse {
status_code: status,
headers: response_headers,
body: body_value,
mode,
};
let behaviors = if latency_ms.is_some() || decorate_fn.is_some() {
let mut behaviors_obj = serde_json::Map::new();
if let Some(ms) = latency_ms {
behaviors_obj.insert("wait".to_string(), serde_json::json!(ms));
}
if let Some(fn_str) = decorate_fn {
behaviors_obj.insert("decorate".to_string(), serde_json::json!(fn_str));
}
Some(serde_json::Value::Object(behaviors_obj))
} else {
None
};
let predicates: Vec<Predicate> = predicates
.into_iter()
.filter_map(|value| match serde_json::from_value(value.clone()) {
Ok(pred) => Some(pred),
Err(e) => {
tracing::warn!(
"Skipping malformed generated predicate: {} (from: {})",
e,
value
);
None
}
})
.collect();
super::types::Stub {
id: None,
predicates,
responses: vec![StubResponse::Is {
is: is_response,
behaviors,
rift: None,
}],
scenario_name: None,
required_scenario_state: None,
new_scenario_state: None,
space: None,
recorded_from,
}
}
pub fn apply_js_or_rhai_decorate(
script: &str,
request: &RequestContext,
body: &str,
status: u16,
headers: &mut HashMap<String, String>,
) -> Result<(String, u16), String> {
if script.trim().starts_with("function") {
#[cfg(feature = "javascript")]
{
let mb_request = crate::scripting::MountebankRequest {
method: request.method.clone(),
path: request.path.clone(),
query: request.query.clone(),
headers: request.headers.clone(),
body: request.body.clone(),
};
match crate::scripting::execute_mountebank_decorate(
script,
&mb_request,
body,
status,
headers,
) {
Ok(result) => {
for (k, v) in result.headers {
headers.insert(k, v);
}
Ok((result.body, result.status_code))
}
Err(e) => Err(format!("JavaScript decorate error: {e}")),
}
}
#[cfg(not(feature = "javascript"))]
{
if let Some(start) = script.find('{') {
if let Some(end) = script.rfind('}') {
let js_body = script[start + 1..end].trim();
let rhai_script = js_body.replace('\'', "\"");
return apply_decorate(&rhai_script, request, body, status, headers);
}
}
Err("Could not parse JavaScript decorate function".to_string())
}
} else {
apply_decorate(script, request, body, status, headers)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_stub_multi_valued_headers_comma_joined() {
let headers = vec![
("Set-Cookie".to_string(), "session=abc".to_string()),
("Set-Cookie".to_string(), "theme=dark".to_string()),
("Content-Type".to_string(), "text/html".to_string()),
];
let stub = create_stub_from_proxy_response(vec![], 200, &headers, b"OK", None, None, None);
match &stub.responses[0] {
StubResponse::Is { is, .. } => {
assert_eq!(
is.headers.get("Set-Cookie").unwrap(),
"session=abc, theme=dark",
"Multi-valued Set-Cookie headers should be comma-joined"
);
assert_eq!(is.headers.get("Content-Type").unwrap(), "text/html");
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_create_stub_hop_by_hop_headers_filtered() {
let headers = vec![
("Content-Type".to_string(), "text/html".to_string()),
("Transfer-Encoding".to_string(), "chunked".to_string()),
("Connection".to_string(), "keep-alive".to_string()),
("Keep-Alive".to_string(), "timeout=5".to_string()),
];
let stub = create_stub_from_proxy_response(vec![], 200, &headers, b"OK", None, None, None);
match &stub.responses[0] {
StubResponse::Is { is, .. } => {
assert!(is.headers.contains_key("Content-Type"));
assert!(
!is.headers.contains_key("Transfer-Encoding"),
"Transfer-Encoding should be filtered"
);
assert!(
!is.headers.contains_key("Connection"),
"Connection should be filtered"
);
assert!(
!is.headers.contains_key("Keep-Alive"),
"Keep-Alive should be filtered"
);
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_create_stub_binary_body_base64_encoded() {
let binary_body: Vec<u8> = vec![0x00, 0xFF, 0xFE, 0xFD, 0x89, 0x50, 0x4E, 0x47];
let stub =
create_stub_from_proxy_response(vec![], 200, &[], &binary_body, None, None, None);
match &stub.responses[0] {
StubResponse::Is { is, .. } => {
assert_eq!(is.mode, ResponseMode::Binary, "Binary body should set mode");
use base64::Engine;
let expected_b64 = base64::engine::general_purpose::STANDARD.encode(&binary_body);
assert_eq!(
is.body.as_ref().unwrap().as_str().unwrap(),
expected_b64,
"Binary body should be base64-encoded"
);
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_create_stub_text_body_not_base64() {
let stub =
create_stub_from_proxy_response(vec![], 200, &[], b"Hello, World!", None, None, None);
match &stub.responses[0] {
StubResponse::Is { is, .. } => {
assert_eq!(
is.mode,
ResponseMode::Text,
"Text body should use text mode"
);
assert_eq!(is.body.as_ref().unwrap().as_str().unwrap(), "Hello, World!");
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_create_stub_json_body_parsed() {
let stub = create_stub_from_proxy_response(
vec![],
200,
&[],
br#"{"key": "value"}"#,
None,
None,
None,
);
match &stub.responses[0] {
StubResponse::Is { is, .. } => {
assert_eq!(is.mode, ResponseMode::Text);
let body = is.body.as_ref().unwrap();
assert!(body.is_object(), "JSON body should be parsed as object");
assert_eq!(body["key"], "value");
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_create_stub_empty_body() {
let stub = create_stub_from_proxy_response(vec![], 204, &[], b"", None, None, None);
match &stub.responses[0] {
StubResponse::Is { is, .. } => {
assert_eq!(is.mode, ResponseMode::Text);
assert!(is.body.is_none(), "Empty body should be None");
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_create_stub_with_latency_and_decorate() {
let stub = create_stub_from_proxy_response(
vec![],
200,
&[],
b"OK",
Some(150),
Some("function(request, response) {}".to_string()),
None,
);
match &stub.responses[0] {
StubResponse::Is { behaviors, .. } => {
let b = behaviors.as_ref().unwrap();
assert_eq!(b["wait"], 150);
assert_eq!(b["decorate"], "function(request, response) {}");
}
_ => panic!("Expected StubResponse::Is"),
}
}
#[test]
fn test_truncate_with_ellipsis_short_string() {
assert_eq!(truncate_with_ellipsis("hello", 10), "hello");
}
#[test]
fn test_truncate_with_ellipsis_exact_length() {
assert_eq!(truncate_with_ellipsis("hello", 5), "hello");
}
#[test]
fn test_truncate_with_ellipsis_long_string() {
assert_eq!(truncate_with_ellipsis("hello world", 5), "hello...");
}
#[test]
fn test_truncate_with_ellipsis_unicode_safe() {
let text = "日本語";
assert_eq!(text.len(), 9);
assert_eq!(truncate_with_ellipsis(text, 5), "日...");
}
#[test]
fn test_truncate_with_ellipsis_emoji() {
let text = "👋🌍🎉";
assert_eq!(truncate_with_ellipsis(text, 5), "👋...");
}
#[test]
fn test_truncate_with_ellipsis_mixed_content() {
let text = "Hello 世界!";
assert_eq!(truncate_with_ellipsis(text, 8), "Hello ...");
}
#[test]
fn test_truncate_with_ellipsis_empty_string() {
assert_eq!(truncate_with_ellipsis("", 10), "");
}
#[test]
fn test_truncate_with_ellipsis_zero_max_len() {
assert_eq!(truncate_with_ellipsis("hello", 0), "...");
}
#[test]
fn test_create_stub_malformed_predicate_skipped() {
let valid_predicate = serde_json::json!({
"equals": { "method": "GET" }
});
let malformed_predicate = serde_json::json!({
"notARealPredicate": { "foo": "bar" }
});
let stub = create_stub_from_proxy_response(
vec![valid_predicate, malformed_predicate],
200,
&[],
b"OK",
None,
None,
None,
);
assert_eq!(stub.predicates.len(), 1);
}
#[test]
fn test_create_stub_all_predicates_malformed() {
let bad1 = serde_json::json!({"garbage": 123});
let bad2 = serde_json::json!("just a string");
let stub =
create_stub_from_proxy_response(vec![bad1, bad2], 200, &[], b"OK", None, None, None);
assert!(stub.predicates.is_empty());
}
#[test]
fn test_create_stub_recorded_from_populated() {
let stub = create_stub_from_proxy_response(
vec![],
200,
&[],
b"OK",
None,
None,
Some("http://upstream:8080".to_string()),
);
assert_eq!(stub.recorded_from.as_deref(), Some("http://upstream:8080"));
}
#[test]
fn test_create_stub_recorded_from_none_when_not_provided() {
let stub = create_stub_from_proxy_response(vec![], 200, &[], b"OK", None, None, None);
assert!(stub.recorded_from.is_none());
}
}