use super::*;
use crate::embed::protocol::EmbedRequest;
use crate::embed::runtime_env::{
DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE, DEFAULT_MIN_AVAILABLE_MB,
DEFAULT_ONNX_INFERENCE_BATCH_SIZE, MAX_ONNX_SEQUENCE_LEN, ONNX_INFERENCE_BATCH_SIZE_ENV,
ONNX_SEQUENCE_LEN_ENV,
};
use std::io::Cursor;
use std::sync::Mutex as StdMutex;
static ENV_LOCK: StdMutex<()> = StdMutex::new(());
struct EnvVarGuard {
key: &'static str,
original: Option<String>,
}
impl EnvVarGuard {
fn set(key: &'static str, value: &str) -> Self {
let original = std::env::var(key).ok();
unsafe {
std::env::set_var(key, value);
}
Self { key, original }
}
fn remove(key: &'static str) -> Self {
let original = std::env::var(key).ok();
unsafe {
std::env::remove_var(key);
}
Self { key, original }
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
match &self.original {
Some(val) => unsafe {
std::env::set_var(self.key, val);
},
None => unsafe {
std::env::remove_var(self.key);
},
}
}
}
fn no_compile_config() -> RuntimeConfig {
RuntimeConfig {
model_name: "__leindex_test_no_model__".to_string(),
rerank_model_name: "__leindex_test_no_rerank_model__".to_string(),
..RuntimeConfig::default()
}
}
#[test]
fn test_runtime_config_default() {
let config = RuntimeConfig::default();
assert_eq!(
config.idle_timeout,
Duration::from_secs(DEFAULT_IDLE_TIMEOUT_SECS)
);
assert_eq!(config.max_frame_size, 16 * 1024 * 1024);
assert_eq!(config.max_text_size, 1024 * 1024);
assert_eq!(config.embedding_dim, 1024);
}
#[test]
fn onnx_inference_batch_size_defaults_to_fixed_batch_safe_value() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
DEFAULT_ONNX_INFERENCE_BATCH_SIZE
);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
1
);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE
);
}
#[test]
fn onnx_inference_batch_size_uses_positive_env_override() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::set(ONNX_INFERENCE_BATCH_SIZE_ENV, "32");
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
32
);
}
#[test]
fn onnx_inference_batch_size_rejects_zero_and_bad_values() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::set(ONNX_INFERENCE_BATCH_SIZE_ENV, "0");
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
DEFAULT_ONNX_INFERENCE_BATCH_SIZE
);
drop(_env);
let _env = EnvVarGuard::set(ONNX_INFERENCE_BATCH_SIZE_ENV, "nope");
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
DEFAULT_ONNX_INFERENCE_BATCH_SIZE
);
}
#[cfg(feature = "onnx")]
#[test]
fn qwen_pooling_uses_last_unpadded_token() {
let runtime = WorkerRuntime::new(no_compile_config());
let pooled = runtime
.pool_and_normalize(
&[
1.0, 0.0, 0.0, 2.0, 8.0, 8.0, ],
1,
3,
&[1, 1, 0],
2,
)
.unwrap();
assert_eq!(pooled.vectors, vec![0.0, 1.0]);
}
#[cfg(feature = "onnx")]
#[test]
fn qwen_pooling_rejects_short_embedding_output() {
let runtime = WorkerRuntime::new(no_compile_config());
let error = runtime
.pool_and_normalize(&[1.0], 1, 2, &[1, 1], 2)
.unwrap_err();
assert_eq!(error.kind, ErrorKind::Inference);
assert!(error.message.contains("embedding output is too short"));
}
#[test]
fn position_ids_repeat_sequence_for_each_batch_row() {
assert_eq!(build_position_ids(2, 4), vec![0, 1, 2, 3, 0, 1, 2, 3]);
}
#[test]
fn onnx_sequence_len_defaults_and_clamps_env_override() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_SEQUENCE_LEN_ENV);
assert_eq!(configured_onnx_sequence_len(), DEFAULT_MAX_SEQ_LEN);
drop(_env);
let _env = EnvVarGuard::set(ONNX_SEQUENCE_LEN_ENV, "4");
assert_eq!(configured_onnx_sequence_len(), DEFAULT_MAX_SEQ_LEN);
drop(_env);
let _env = EnvVarGuard::set(ONNX_SEQUENCE_LEN_ENV, "256");
assert_eq!(configured_onnx_sequence_len(), 256);
drop(_env);
let _env = EnvVarGuard::set(ONNX_SEQUENCE_LEN_ENV, "4096");
assert_eq!(configured_onnx_sequence_len(), MAX_ONNX_SEQUENCE_LEN);
}
#[test]
fn dynamic_qwen_uses_batched_inference_by_default() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "cpu"),
DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE
);
const _: () = assert!(DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE > 1);
}
#[test]
fn migraphx_uses_one_stable_batch_shape_by_default() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "migraphx"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE
);
}
#[test]
fn test_batch_size_for_dynamic_uint8() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic-uint8", "migraphx"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE
);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic-uint8", "cpu"),
DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE
);
}
#[cfg(feature = "onnx")]
#[test]
fn test_extract_u8_dequantization() {
const SCALE: f32 = 0.002_745_098;
const ZERO_POINT: f32 = 109.0;
let test_cases: [(u8, f32); 4] = [
(109, 0.0), (0, (0.0 - 109.0) * SCALE), (255, (255.0 - 109.0) * SCALE), (128, (128.0 - 109.0) * SCALE), ];
for (input, expected) in test_cases {
let dequantized = (input as f32 - ZERO_POINT) * SCALE;
assert!(
(dequantized - expected).abs() < 1e-6,
"u8 value {}: expected {}, got {}",
input,
expected,
dequantized
);
}
}
#[cfg(feature = "onnx")]
#[test]
fn test_u8_dequant_preserves_unit_norm() {
const SCALE: f32 = 0.002_745_098;
const ZERO_POINT: f32 = 109.0;
let original: Vec<f32> = vec![0.35; 8];
let quantized: Vec<u8> = original
.iter()
.map(|&v| {
let q = (v / SCALE + ZERO_POINT).round() as i32;
q.clamp(0, 255) as u8
})
.collect();
let dequantized: Vec<f32> = quantized
.iter()
.map(|&v| (v as f32 - ZERO_POINT) * SCALE)
.collect();
let norm: f32 = dequantized.iter().map(|v| v * v).sum::<f32>().sqrt();
assert!(
(norm - 1.0).abs() < 0.1,
"dequantized vector norm {} should be close to 1.0",
norm
);
}
#[test]
fn test_runtime_idle_not_expired_initially() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
assert!(!rt.is_idle_expired());
}
#[test]
fn test_runtime_idle_expired_with_zero_timeout() {
let config = RuntimeConfig {
idle_timeout: Duration::from_secs(0),
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
std::thread::sleep(Duration::from_millis(1));
assert!(rt.is_idle_expired());
}
#[test]
fn test_runtime_touch_resets_idle() {
let config = RuntimeConfig {
idle_timeout: Duration::from_millis(10),
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
std::thread::sleep(Duration::from_millis(20));
assert!(rt.is_idle_expired());
rt.touch();
assert!(!rt.is_idle_expired());
}
#[test]
fn cloned_runtime_shares_idle_activity() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<WorkerRuntime>();
let runtime = WorkerRuntime::new(no_compile_config());
let cloned = runtime.clone();
assert!(Arc::ptr_eq(&runtime.last_activity, &cloned.last_activity));
cloned.touch();
assert!(!runtime.is_idle_expired());
}
#[test]
fn test_shutdown_flag() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let flag = rt.shutdown_flag();
assert!(!flag.load(Ordering::Relaxed));
flag.store(true, Ordering::Relaxed);
assert!(flag.load(Ordering::Relaxed));
}
#[test]
fn test_truncate_text_within_limit() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let text = "hello world".to_string();
let result = rt.truncate_text(text.clone());
assert_eq!(result, text);
}
#[test]
fn test_truncate_text_exceeds_limit() {
let config = RuntimeConfig {
max_text_size: 10,
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
let text = "hello world, this is a long string".to_string();
let result = rt.truncate_text(text);
assert!(result.len() <= 10);
assert_eq!(result, "hello worl");
}
#[test]
fn test_truncate_text_unicode_boundary() {
let config = RuntimeConfig {
max_text_size: 10,
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
let text = "héllo wörld test".to_string();
let result = rt.truncate_text(text);
assert!(result.len() <= 10);
assert!(result.is_char_boundary(result.len()));
}
#[test]
fn test_handle_embed_empty_batch() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let request = EmbedRequest {
texts: vec![],
expected_dim: 1024,
cache_keys: vec![],
};
let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
let result = rt.handle_embed(&frame, &Arc::new(AtomicBool::new(false)));
let response = result.unwrap();
assert_eq!(response.count, 0);
assert_eq!(response.dimension, 1024);
assert!(response.vectors.is_empty());
}
#[test]
fn test_handle_embed_returns_flat_row_major() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let request = EmbedRequest {
texts: vec!["hello".to_string(), "world".to_string()],
expected_dim: 8,
cache_keys: vec![],
};
let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
let result = rt.handle_embed(&frame, &Arc::new(AtomicBool::new(false)));
#[cfg(feature = "onnx")]
{
let err = result.unwrap_err();
assert!(
err.kind == ErrorKind::ModelNotFound || err.kind == ErrorKind::Inference,
"expected ModelNotFound or Inference, got {:?}: {}",
err.kind,
err.message
);
}
#[cfg(not(feature = "onnx"))]
{
let response = result.unwrap();
assert_eq!(response.count, 2);
assert_eq!(response.dimension, 8);
assert_eq!(response.vectors.len(), 16);
assert_eq!(response.get_embedding(0).unwrap().len(), 8);
assert_eq!(response.get_embedding(1).unwrap().len(), 8);
}
}
#[test]
fn test_handle_embed_preserves_ordering() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let texts: Vec<String> = (0..5).map(|i| format!("text {}", i)).collect();
let request = EmbedRequest {
texts: texts.clone(),
expected_dim: 4,
cache_keys: vec![],
};
let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
let result = rt.handle_embed(&frame, &Arc::new(AtomicBool::new(false)));
#[cfg(feature = "onnx")]
{
let err = result.unwrap_err();
assert!(
err.kind == ErrorKind::ModelNotFound || err.kind == ErrorKind::Inference,
"expected ModelNotFound or Inference, got {:?}: {}",
err.kind,
err.message
);
}
#[cfg(not(feature = "onnx"))]
{
let response = result.unwrap();
assert_eq!(response.count, 5);
for i in 0..5 {
assert!(response.get_embedding(i).is_some());
}
}
}
#[test]
fn test_dispatch_embed_request() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let request = EmbedRequest {
texts: vec!["test".to_string()],
expected_dim: 4,
cache_keys: vec![],
};
let frame = protocol::embed_request_frame(BatchId::new(42), request).unwrap();
let response_frame = rt.dispatch(&frame);
assert_eq!(response_frame.header.batch_id, BatchId::new(42));
#[cfg(feature = "onnx")]
{
assert_eq!(response_frame.header.msg_type, MsgType::Error);
}
#[cfg(not(feature = "onnx"))]
{
assert_eq!(response_frame.header.msg_type, MsgType::EmbedResponse);
}
}
#[test]
fn test_dispatch_rerank_request() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let request = protocol::RerankRequest {
query: "test".to_string(),
documents: vec![protocol::RerankDocument {
id: "doc1".to_string(),
content: "content".to_string(),
initial_score: 0.9,
}],
};
let frame = protocol::rerank_request_frame(BatchId::new(7), request).unwrap();
let response_frame = rt.dispatch(&frame);
assert_eq!(response_frame.header.batch_id, BatchId::new(7));
#[cfg(feature = "onnx")]
{
assert_eq!(response_frame.header.msg_type, MsgType::Error);
}
#[cfg(not(feature = "onnx"))]
{
assert_eq!(response_frame.header.msg_type, MsgType::RerankResponse);
}
}
#[test]
fn test_dispatch_unknown_message_type() {
let config = no_compile_config();
let rt = WorkerRuntime::new(config);
let frame = Frame {
header: protocol::FrameHeader {
batch_id: BatchId::new(99),
msg_type: MsgType::Error, },
payload: vec![],
};
let response_frame = rt.dispatch(&frame);
assert_eq!(response_frame.header.batch_id, BatchId::new(99));
assert_eq!(response_frame.header.msg_type, MsgType::Error);
}
#[test]
fn test_run_loop_single_request() {
let config = RuntimeConfig {
idle_timeout: Duration::from_secs(300),
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
let request = EmbedRequest {
texts: vec!["hello".to_string()],
expected_dim: 4,
cache_keys: vec![],
};
let frame = protocol::embed_request_frame(BatchId::new(1), request).unwrap();
let wire = frame.encode_wire().unwrap();
let reader = Cursor::new(wire);
let writer = Cursor::new(Vec::<u8>::new());
let result = rt.run_loop(reader, writer);
assert!(result.is_ok());
}
#[test]
fn runtime_config_provider_env_overrides_toml_and_default() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::set("LEINDEX_WORKER_EXECUTION_PROVIDER", "cuda");
let config = RuntimeConfig::from_env();
assert_eq!(
config.execution_provider, "cuda",
"env var must take precedence over TOML and the 'auto' default"
);
drop(_env);
let _env_removed = EnvVarGuard::remove("LEINDEX_WORKER_EXECUTION_PROVIDER");
let config = RuntimeConfig::from_env();
assert!(
!config.execution_provider.trim().is_empty(),
"provider must always resolve to a non-empty value"
);
assert_eq!(
config.execution_provider,
config.execution_provider.to_ascii_lowercase(),
);
}
#[test]
fn runtime_config_provider_blanks_env_falls_through() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::set("LEINDEX_WORKER_EXECUTION_PROVIDER", " ");
let config = RuntimeConfig::from_env();
assert!(
!config.execution_provider.trim().is_empty(),
"blank env value must fall through to a non-empty default"
);
}
#[test]
fn runtime_config_min_available_mb_defaults_to_documented_floor() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove("LEINDEX_WORKER_MIN_AVAILABLE_MB");
let config = RuntimeConfig::from_env();
assert_eq!(
config.min_available_mb,
Some(DEFAULT_MIN_AVAILABLE_MB),
"unset MIN_AVAILABLE_MB must apply the documented 2048 MiB default"
);
let _env2 = EnvVarGuard::set("LEINDEX_WORKER_MIN_AVAILABLE_MB", "4096");
let config = RuntimeConfig::from_env();
assert_eq!(config.min_available_mb, Some(4096));
let _env3 = EnvVarGuard::set("LEINDEX_WORKER_MIN_AVAILABLE_MB", "0");
let config = RuntimeConfig::from_env();
assert_eq!(config.min_available_mb, None, "0 = disabled");
let _env4 = EnvVarGuard::set("LEINDEX_WORKER_MIN_AVAILABLE_MB", "not-a-number");
let config = RuntimeConfig::from_env();
assert_eq!(
config.min_available_mb,
Some(DEFAULT_MIN_AVAILABLE_MB),
"malformed MIN_AVAILABLE_MB must keep the documented 2048 MiB default"
);
}
#[test]
fn explicit_gpu_provider_select_returns_cpu_fallback_not_hard_error() {
use crate::embed::provider::ExecutionProviderSelector;
let result = ExecutionProviderSelector::select("cuda");
match result {
Ok(selection) => {
let name = selection.name();
assert!(
name == "cuda" || name == "cpu",
"unexpected cuda-resolution: {name}"
);
}
Err(fallback) => {
assert_eq!(fallback.fallback_name(), "cpu");
assert!(!fallback.is_requested_provider());
assert!(
fallback.reason().contains("CUDA"),
"fallback reason should name the missing provider: {}",
fallback.reason()
);
}
}
}
#[cfg(feature = "onnx")]
#[test]
fn runtime_resolved_provider_is_never_auto_for_embed_or_rerank() {
let rt = WorkerRuntime::new(RuntimeConfig {
execution_provider: "auto".to_string(),
..no_compile_config()
});
let health = rt.health_response(crate::embed::protocol::WorkerState::Initializing, None);
let provider = health.provider.expect("health response carries a provider");
assert_ne!(
provider, "auto",
"the provider handed to embedder/reranker session builders must be concrete"
);
assert!(
matches!(provider.as_str(), "cpu" | "cuda" | "migraphx" | "coreml"),
"resolved provider must be one of the concrete tokens, got {provider}"
);
}
#[test]
fn test_run_loop_multiple_requests_same_runtime() {
let config = RuntimeConfig {
idle_timeout: Duration::from_secs(300),
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
let request1 = EmbedRequest {
texts: vec!["first".to_string()],
expected_dim: 4,
cache_keys: vec![],
};
let request2 = EmbedRequest {
texts: vec!["second".to_string()],
expected_dim: 4,
cache_keys: vec![],
};
let frame1 = protocol::embed_request_frame(BatchId::new(1), request1).unwrap();
let frame2 = protocol::embed_request_frame(BatchId::new(2), request2).unwrap();
let wire1 = frame1.encode_wire().unwrap();
let wire2 = frame2.encode_wire().unwrap();
let mut combined = wire1.clone();
combined.extend_from_slice(&wire2);
let reader = Cursor::new(combined);
let writer = Cursor::new(Vec::<u8>::new());
let result = rt.run_loop(reader, writer);
assert!(result.is_ok());
result.unwrap();
}
#[test]
fn test_idle_timeout_causes_exit() {
let config = RuntimeConfig {
idle_timeout: Duration::from_millis(1),
..no_compile_config()
};
let rt = WorkerRuntime::new(config);
let reader = Cursor::new(Vec::<u8>::new());
let writer = Cursor::new(Vec::<u8>::new());
let result = rt.run_loop(reader, writer);
assert!(result.is_ok());
}
#[cfg(feature = "onnx")]
fn test_encoding(marker: u64) -> tokenizers::Encoding {
use std::collections::HashMap;
tokenizers::Encoding::new(
vec![marker as u32, (marker + 1) as u32],
vec![0, 0],
vec![marker.to_string(), (marker + 1).to_string()],
vec![None, None],
vec![(0, 1), (1, 2)],
vec![0, 0],
vec![1, 1],
vec![],
HashMap::new(),
)
}
#[cfg(feature = "onnx")]
fn embed_encodings(n: usize) -> Vec<tokenizers::Encoding> {
(0..n).map(|i| test_encoding(i as u64)).collect()
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_batch_cpu_all_sub_batches_processed() {
let rt = WorkerRuntime::new(no_compile_config());
let encodings = embed_encodings(17);
let batch_size = 8usize;
let dim = 4usize;
let mut batches: Vec<usize> = Vec::new();
let all_pooled = rt
.run_onnx_embed_batch_loop(&encodings, batch_size, false, dim, |sub, dim| {
batches.push(sub.len());
let mut out = Vec::with_capacity(sub.len() * dim);
for (i, _) in sub.iter().enumerate() {
out.extend(std::iter::repeat_n(i as f32, dim));
}
Ok(out)
})
.unwrap();
assert_eq!(batches, vec![8, 8, 1]);
assert_eq!(all_pooled.len(), 17 * dim);
assert_ne!(all_pooled.len(), 0, "CPU/CUDA embed must not be empty");
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_batch_migraphx_full_sub_batches_no_padding() {
let rt = WorkerRuntime::new(no_compile_config());
let batch_size = 8usize;
let dim = 4usize;
let encodings = embed_encodings(16);
let mut batches: Vec<usize> = Vec::new();
let all_pooled = rt
.run_onnx_embed_batch_loop(&encodings, batch_size, true, dim, |sub, dim| {
batches.push(sub.len());
Ok(vec![1.0f32; sub.len() * dim])
})
.unwrap();
assert_eq!(batches, vec![8, 8]);
assert_eq!(batches.len(), 2);
assert_eq!(all_pooled.len(), 16 * dim);
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_batch_migraphx_partial_sub_batch_padded_and_trimmed() {
let rt = WorkerRuntime::new(no_compile_config());
let batch_size = 8usize;
let dim = 4usize;
let encodings = embed_encodings(10);
let mut batches: Vec<usize> = Vec::new();
let all_pooled = rt
.run_onnx_embed_batch_loop(&encodings, batch_size, true, dim, |sub, dim| {
batches.push(sub.len());
let mut out = Vec::with_capacity(sub.len() * dim);
for (i, _) in sub.iter().enumerate() {
out.extend(std::iter::repeat_n(i as f32, dim));
}
Ok(out)
})
.unwrap();
assert_eq!(batches, vec![8, 8]);
assert_eq!(all_pooled.len(), 10 * dim);
let tail = &all_pooled[all_pooled.len() - dim..];
assert_eq!(
tail, &[1.0f32; 4],
"last row must be the real row 1, not padding"
);
assert_eq!(*all_pooled.last().unwrap(), 1.0);
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_response_invariant_all_providers() {
let rt = WorkerRuntime::new(no_compile_config());
let batch_size = 8usize;
let dim = 4usize;
let providers: &[(&str, bool)] = &[("cpu", false), ("migraphx", true)];
let counts: &[usize] = &[1, 3, 8, 9, 16, 17];
for (provider, fixed_batch) in providers {
for &count in counts {
let encodings = embed_encodings(count);
let mut batches: Vec<usize> = Vec::new();
let all_pooled = rt
.run_onnx_embed_batch_loop(&encodings, batch_size, *fixed_batch, dim, |sub, dim| {
batches.push(sub.len());
Ok(vec![0.5f32; sub.len() * dim])
})
.unwrap();
assert_eq!(
all_pooled.len(),
count * dim,
"{provider} count={count}: vectors.len() != count*dimension"
);
let response = EmbedResponse::new(all_pooled.clone(), count, dim);
assert_eq!(response.vectors.len(), count * dim);
assert_eq!(response.count, count);
assert_eq!(response.dimension, dim);
assert_eq!(
response.vectors.len(),
response.count * response.dimension,
"EmbedResponse invariant broken for {provider} count={count}"
);
}
}
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_tokenizes_and_infers_per_sub_batch() {
let rt = WorkerRuntime::new(no_compile_config());
let texts: Vec<String> = (0..17).map(|i| format!("text-{i}")).collect();
let mut tokenized_sizes = Vec::new();
let events = std::cell::RefCell::new(Vec::new());
let dim = 4usize;
let pooled = rt
.run_onnx_embed_text_batch_loop(
&texts,
8,
false,
dim,
&Arc::new(AtomicBool::new(false)),
|sub_texts| {
tokenized_sizes.push(sub_texts.len());
events
.borrow_mut()
.push(format!("tokenize-{}", sub_texts.len()));
Ok(embed_encodings(sub_texts.len()))
},
|encodings, dim| {
events
.borrow_mut()
.push(format!("infer-{}", encodings.len()));
Ok(vec![1.0f32; encodings.len() * dim])
},
)
.unwrap();
assert_eq!(tokenized_sizes, vec![8, 8, 1]);
assert_eq!(pooled.len(), texts.len() * dim);
assert_eq!(
events.into_inner(),
vec![
"tokenize-8".to_string(),
"infer-8".to_string(),
"tokenize-8".to_string(),
"infer-8".to_string(),
"tokenize-1".to_string(),
"infer-1".to_string(),
]
);
}
#[cfg(feature = "onnx")]
#[test]
fn test_cpu_only_fallback_preserves_vectors_across_batch_sizes() {
use crate::embed::provider::{Provider, select_auto_from_availability};
let provider = select_auto_from_availability(false, false, false);
assert_eq!(provider, Provider::Cpu);
let rt = WorkerRuntime::new(RuntimeConfig {
execution_provider: provider.name(),
..no_compile_config()
});
let texts: Vec<String> = (0..17).map(|i| i.to_string()).collect();
let expected: Vec<f32> = (0..17).flat_map(|i| [i as f32, -(i as f32)]).collect();
for batch_size in [1, 8, 32] {
let actual = rt
.run_onnx_embed_text_batch_loop(
&texts,
batch_size,
false,
2,
&Arc::new(AtomicBool::new(false)),
|texts| Ok(texts.to_vec()),
|texts, _| {
Ok(texts
.iter()
.flat_map(|text| {
let value = text.parse::<f32>().unwrap();
[value, -value]
})
.collect())
},
)
.unwrap();
assert_eq!(
actual, expected,
"CPU fallback changed rows at batch {batch_size}"
);
}
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_fixed_batch_pads_after_per_batch_tokenization() {
let rt = WorkerRuntime::new(no_compile_config());
let texts: Vec<String> = (0..10).map(|i| format!("text-{i}")).collect();
let mut tokenized_sizes = Vec::new();
let mut inferred_sizes = Vec::new();
let dim = 2usize;
let pooled = rt
.run_onnx_embed_text_batch_loop(
&texts,
8,
true,
dim,
&Arc::new(AtomicBool::new(false)),
|sub_texts| {
tokenized_sizes.push(sub_texts.len());
Ok(embed_encodings(sub_texts.len()))
},
|encodings, dim| {
inferred_sizes.push(encodings.len());
Ok(vec![0.5f32; encodings.len() * dim])
},
)
.unwrap();
assert_eq!(tokenized_sizes, vec![8, 2]);
assert_eq!(inferred_sizes, vec![8, 8]);
assert_eq!(pooled.len(), texts.len() * dim);
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_later_tokenizer_error_stops_before_next_inference() {
let rt = WorkerRuntime::new(no_compile_config());
let texts: Vec<String> = (0..10).map(|i| format!("text-{i}")).collect();
let mut inferred_batches = Vec::new();
let mut tokenizer_calls = 0usize;
let error = rt
.run_onnx_embed_text_batch_loop(
&texts,
8,
false,
2,
&Arc::new(AtomicBool::new(false)),
|_sub_texts| {
tokenizer_calls += 1;
if tokenizer_calls == 2 {
Err(WorkerError {
kind: ErrorKind::Tokenizer,
message: "synthetic tokenizer failure".to_string(),
})
} else {
Ok(embed_encodings(8))
}
},
|encodings, dim| {
inferred_batches.push(encodings.len());
Ok(vec![1.0f32; encodings.len() * dim])
},
)
.unwrap_err();
assert_eq!(inferred_batches, vec![8]);
assert_eq!(tokenizer_calls, 2);
assert_eq!(error.kind, ErrorKind::Tokenizer);
assert!(error.message.contains("synthetic tokenizer failure"));
}
#[cfg(feature = "onnx")]
#[test]
fn test_embed_text_batch_loop_checks_cancel_between_sub_batches() {
let rt = WorkerRuntime::new(no_compile_config());
let texts: Vec<String> = (0..10).map(|i| format!("text-{i}")).collect();
let mut inferred_batches = Vec::new();
let cancel_token = Arc::new(AtomicBool::new(false));
let cancel_from_runner = Arc::clone(&cancel_token);
let error = rt
.run_onnx_embed_text_batch_loop(
&texts,
8,
false,
2,
&cancel_token,
|_sub_texts| Ok(embed_encodings(8)),
|encodings, dim| {
inferred_batches.push(encodings.len());
cancel_from_runner.store(true, Ordering::Release);
Ok(vec![1.0f32; encodings.len() * dim])
},
)
.unwrap_err();
assert_eq!(inferred_batches, vec![8]);
assert_eq!(error.kind, ErrorKind::Inference);
assert!(error.message.contains("cancelled"));
}
#[test]
fn test_non_dynamic_model_migraphx_returns_fixed_batch() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
"non-dynamic model with migraphx must return the fixed MIGraphX batch (the static graph is b8-s128)"
);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "migraphx"),
8
);
}
#[test]
fn test_non_dynamic_model_rocm_returns_fixed_batch() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "rocm"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
"non-dynamic model with rocm must return the fixed MIGraphX batch"
);
}
#[test]
fn test_dynamic_model_migraphx_returns_batch_8() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "migraphx"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
"dynamic model with migraphx must return the MIGraphX stable batch size"
);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "migraphx"),
8
);
}
#[test]
fn test_dynamic_model_rocm_returns_batch_8() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "rocm"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
"dynamic model with rocm must return the MIGraphX stable batch size"
);
}
#[test]
fn test_non_dynamic_model_cpu_returns_batch_1() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "cpu"),
DEFAULT_ONNX_INFERENCE_BATCH_SIZE,
"non-dynamic model with cpu must return batch_size=1 (unchanged)"
);
}
#[test]
fn test_dynamic_model_cpu_returns_default_dynamic() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b-dynamic", "cpu"),
DEFAULT_DYNAMIC_ONNX_INFERENCE_BATCH_SIZE,
"dynamic model with cpu must return the dynamic batch size"
);
}
#[test]
fn test_non_dynamic_model_migraphx_case_insensitive() {
let _guard = ENV_LOCK.lock().unwrap_or_else(|p| p.into_inner());
let _env = EnvVarGuard::remove(ONNX_INFERENCE_BATCH_SIZE_ENV);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "MIGRAPHX"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
"non-dynamic model with uppercase MIGRAPHX must return the fixed MIGraphX batch"
);
assert_eq!(
configured_onnx_inference_batch_size("qwen3-embed-0.6b", "ROCm"),
DEFAULT_MIGRAPHX_INFERENCE_BATCH_SIZE,
);
}
#[cfg(feature = "onnx")]
#[test]
fn test_collapsed_batch_sentinel_is_detectable() {
assert!(COLLAPSED_BATCH_SENTINEL.starts_with("__"));
assert!(!COLLAPSED_BATCH_SENTINEL.is_empty());
}
#[cfg(feature = "onnx")]
#[test]
fn test_finalize_embed_output_detects_collapsed_batch() {
let fake_error_msg = format!(
"{}: model collapsed batch dimension (sent 3, got [1, 128, 1024])",
COLLAPSED_BATCH_SENTINEL
);
assert!(
fake_error_msg.starts_with(COLLAPSED_BATCH_SENTINEL),
"collapsed batch error must start with the sentinel for retry detection"
);
}
#[cfg(feature = "onnx")]
#[test]
fn test_non_collapsed_error_does_not_match_sentinel() {
let regular_error = "ONNX inference failed: shape mismatch";
assert!(
!regular_error.starts_with(COLLAPSED_BATCH_SENTINEL),
"regular errors must not match the collapsed batch sentinel"
);
}