use super::compute::{
resolve_coreml_batch_size, resolve_coreml_tripwire_mb, COREML_BATCH_SIZE_MAX,
DEFAULT_COREML_BATCH_SIZE, DEFAULT_COREML_TRIPWIRE_MB,
};
use super::detect::detect_total_ram_mb;
use super::policy::MemoryPolicy;
use super::tier::MemoryTier;
use std::sync::Mutex;
pub(super) static ENV_LOCK: Mutex<()> = Mutex::new(());
#[test]
fn test_index_memory_limit_resolution_floors_at_global() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior_idx = std::env::var("TRUSTY_INDEX_MEMORY_LIMIT_MB").ok();
let prior_mem = std::env::var("TRUSTY_MEMORY_LIMIT_MB").ok();
unsafe {
std::env::set_var("TRUSTY_MEMORY_LIMIT_MB", "8192");
std::env::set_var("TRUSTY_INDEX_MEMORY_LIMIT_MB", "2048"); }
let policy = MemoryPolicy::from_total_ram_mb(32 * 1024);
assert_eq!(policy.memory_limit_mb, 8_192);
assert_eq!(policy.index_memory_limit_mb, 8_192);
unsafe {
match prior_idx {
Some(v) => std::env::set_var("TRUSTY_INDEX_MEMORY_LIMIT_MB", v),
None => std::env::remove_var("TRUSTY_INDEX_MEMORY_LIMIT_MB"),
}
match prior_mem {
Some(v) => std::env::set_var("TRUSTY_MEMORY_LIMIT_MB", v),
None => std::env::remove_var("TRUSTY_MEMORY_LIMIT_MB"),
}
}
}
#[test]
fn test_memory_limit_scales_proportionally_across_xlarge_hosts() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_MEMORY_LIMIT_MB").ok();
unsafe {
std::env::remove_var("TRUSTY_MEMORY_LIMIT_MB");
}
let p64 = MemoryPolicy::from_total_ram_mb(64 * 1024);
unsafe {
std::env::remove_var("TRUSTY_MEMORY_LIMIT_MB");
}
let p128 = MemoryPolicy::from_total_ram_mb(128 * 1024);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_MEMORY_LIMIT_MB", v),
None => std::env::remove_var("TRUSTY_MEMORY_LIMIT_MB"),
}
}
assert_eq!(p64.tier, MemoryTier::XLarge);
assert_eq!(p128.tier, MemoryTier::XLarge);
assert!(
p128.memory_limit_mb > p64.memory_limit_mb,
"128 GB host ({} MB) should have a larger memory_limit_mb than \
a 64 GB host ({} MB) — see issue #120",
p128.memory_limit_mb,
p64.memory_limit_mb,
);
assert_eq!(p64.memory_limit_mb, 16 * 1024);
assert_eq!(p128.memory_limit_mb, 32 * 1024);
}
#[test]
fn test_env_override() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_MAX_CHUNKS").ok();
unsafe {
std::env::set_var("TRUSTY_MAX_CHUNKS", "42");
}
let policy = MemoryPolicy::from_total_ram_mb(16 * 1024);
assert_eq!(policy.tier, MemoryTier::Medium);
assert_eq!(policy.max_chunks, 42);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_MAX_CHUNKS", v),
None => std::env::remove_var("TRUSTY_MAX_CHUNKS"),
}
}
}
#[test]
fn test_tier_batch_size_hard_cap() {
assert_eq!(MemoryTier::Medium.batch_size_hard_cap(), 128);
assert_eq!(MemoryTier::Large.batch_size_hard_cap(), 256);
assert_eq!(MemoryTier::XLarge.batch_size_hard_cap(), 512);
}
#[test]
fn test_batch_size_env_override_clamped_by_hard_cap() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_MAX_BATCH_SIZE").ok();
let prior_explicit = std::env::var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT").ok();
unsafe {
std::env::set_var("TRUSTY_MAX_BATCH_SIZE", "2048");
std::env::remove_var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT");
}
let policy = MemoryPolicy::from_total_ram_mb(16 * 1024);
assert_eq!(policy.tier, MemoryTier::Medium);
assert_eq!(
policy.max_batch_size, 128,
"Medium tier must clamp TRUSTY_MAX_BATCH_SIZE=2048 down to 128"
);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_MAX_BATCH_SIZE", v),
None => std::env::remove_var("TRUSTY_MAX_BATCH_SIZE"),
}
match prior_explicit {
Some(v) => std::env::set_var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT", v),
None => std::env::remove_var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT"),
}
}
}
#[test]
fn test_batch_size_explicit_flag_bypasses_clamp() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_MAX_BATCH_SIZE").ok();
let prior_explicit = std::env::var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT").ok();
let prior_mem = std::env::var("TRUSTY_MEMORY_LIMIT_MB").ok();
unsafe {
std::env::set_var("TRUSTY_MAX_BATCH_SIZE", "512");
std::env::set_var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT", "1");
std::env::remove_var("TRUSTY_MEMORY_LIMIT_MB");
}
let policy = MemoryPolicy::from_total_ram_mb(16 * 1024);
assert_eq!(policy.tier, MemoryTier::Medium);
assert_eq!(
policy.max_batch_size, 512,
"TRUSTY_MAX_BATCH_SIZE_EXPLICIT=1 must bypass the tier hard cap"
);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_MAX_BATCH_SIZE", v),
None => std::env::remove_var("TRUSTY_MAX_BATCH_SIZE"),
}
match prior_explicit {
Some(v) => std::env::set_var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT", v),
None => std::env::remove_var("TRUSTY_MAX_BATCH_SIZE_EXPLICIT"),
}
match prior_mem {
Some(v) => std::env::set_var("TRUSTY_MEMORY_LIMIT_MB", v),
None => std::env::remove_var("TRUSTY_MEMORY_LIMIT_MB"),
}
}
}
#[test]
fn test_coreml_batch_size_default() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_COREML_BATCH_SIZE").ok();
unsafe { std::env::remove_var("TRUSTY_COREML_BATCH_SIZE") };
assert_eq!(resolve_coreml_batch_size(), DEFAULT_COREML_BATCH_SIZE);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_COREML_BATCH_SIZE", v),
None => std::env::remove_var("TRUSTY_COREML_BATCH_SIZE"),
}
}
}
#[test]
fn test_coreml_batch_size_env_override() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_COREML_BATCH_SIZE").ok();
unsafe { std::env::set_var("TRUSTY_COREML_BATCH_SIZE", "64") };
assert_eq!(resolve_coreml_batch_size(), 64);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_COREML_BATCH_SIZE", v),
None => std::env::remove_var("TRUSTY_COREML_BATCH_SIZE"),
}
}
}
#[test]
fn test_coreml_batch_size_env_clamp() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_COREML_BATCH_SIZE").ok();
unsafe { std::env::set_var("TRUSTY_COREML_BATCH_SIZE", "10000") };
assert_eq!(resolve_coreml_batch_size(), COREML_BATCH_SIZE_MAX);
unsafe { std::env::set_var("TRUSTY_COREML_BATCH_SIZE", "0") };
assert_eq!(resolve_coreml_batch_size(), DEFAULT_COREML_BATCH_SIZE);
unsafe { std::env::set_var("TRUSTY_COREML_BATCH_SIZE", "not-a-number") };
assert_eq!(resolve_coreml_batch_size(), DEFAULT_COREML_BATCH_SIZE);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_COREML_BATCH_SIZE", v),
None => std::env::remove_var("TRUSTY_COREML_BATCH_SIZE"),
}
}
}
#[test]
fn test_coreml_tripwire_default() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_COREML_TRIPWIRE_MB").ok();
unsafe { std::env::remove_var("TRUSTY_COREML_TRIPWIRE_MB") };
assert_eq!(resolve_coreml_tripwire_mb(), DEFAULT_COREML_TRIPWIRE_MB);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_COREML_TRIPWIRE_MB", v),
None => std::env::remove_var("TRUSTY_COREML_TRIPWIRE_MB"),
}
}
}
#[test]
fn test_coreml_tripwire_env_override() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_COREML_TRIPWIRE_MB").ok();
unsafe { std::env::set_var("TRUSTY_COREML_TRIPWIRE_MB", "8192") };
assert_eq!(resolve_coreml_tripwire_mb(), 8192);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_COREML_TRIPWIRE_MB", v),
None => std::env::remove_var("TRUSTY_COREML_TRIPWIRE_MB"),
}
}
}
#[test]
fn test_coreml_tripwire_env_invalid() {
let _g = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let prior = std::env::var("TRUSTY_COREML_TRIPWIRE_MB").ok();
unsafe { std::env::set_var("TRUSTY_COREML_TRIPWIRE_MB", "0") };
assert_eq!(resolve_coreml_tripwire_mb(), DEFAULT_COREML_TRIPWIRE_MB);
unsafe { std::env::set_var("TRUSTY_COREML_TRIPWIRE_MB", "not-a-number") };
assert_eq!(resolve_coreml_tripwire_mb(), DEFAULT_COREML_TRIPWIRE_MB);
unsafe {
match prior {
Some(v) => std::env::set_var("TRUSTY_COREML_TRIPWIRE_MB", v),
None => std::env::remove_var("TRUSTY_COREML_TRIPWIRE_MB"),
}
}
}
#[test]
fn test_ram_detection_returns_nonzero() {
if let Some(mb) = detect_total_ram_mb() {
assert!(mb > 0, "detected RAM should be > 0, got {mb}");
assert!(
mb < 4 * 1024 * 1024,
"detected RAM implausibly large: {mb} MB"
);
}
}