[model-gateway] : Rust integration tests for integration_mock replacement (#16441)

This commit is contained in:
Simo Lin
2026-01-04 21:01:47 -08:00
committed by GitHub
parent 078270473a
commit f84487af59
48 changed files with 4856 additions and 57 deletions
@@ -23,6 +23,12 @@ pub struct MockSearchServer {
tool_router: ToolRouter<MockSearchServer>,
}
impl Default for MockSearchServer {
fn default() -> Self {
Self::new()
}
}
#[tool_router]
impl MockSearchServer {
pub fn new() -> Self {
+6
View File
@@ -6,7 +6,11 @@ pub mod mock_openai_server;
pub mod mock_worker;
pub mod streaming_helpers;
pub mod test_app;
pub mod test_certs;
pub mod test_config;
pub mod tls_mock_worker;
// Re-export commonly used test builders
use std::{
fs,
path::PathBuf,
@@ -33,6 +37,8 @@ use smg::{
tokenizer::registry::TokenizerRegistry,
tool_parser::ParserFactory as ToolParserFactory,
};
#[allow(unused_imports)]
pub use test_config::{TestRouterConfig, TestWorkerConfig};
/// Test context for directly testing mock workers without full router setup.
pub struct WorkerTestContext {
@@ -0,0 +1,299 @@
// TLS certificate generation for integration tests
#![allow(dead_code)]
use std::path::PathBuf;
use openssl::{
asn1::Asn1Time,
bn::{BigNum, MsbOption},
hash::MessageDigest,
pkey::{PKey, Private},
rsa::Rsa,
x509::{
extension::{BasicConstraints, ExtendedKeyUsage, KeyUsage, SubjectAlternativeName},
X509Builder, X509NameBuilder, X509,
},
};
use tempfile::TempDir;
/// Container for generated test certificates
pub struct TestCertificates {
/// Temporary directory containing the certificate files
pub temp_dir: TempDir,
/// Path to CA certificate
pub ca_cert_path: PathBuf,
/// Path to CA private key
pub ca_key_path: PathBuf,
/// Path to server certificate
pub server_cert_path: PathBuf,
/// Path to server private key
pub server_key_path: PathBuf,
/// Path to client certificate
pub client_cert_path: PathBuf,
/// Path to client private key
pub client_key_path: PathBuf,
}
impl TestCertificates {
/// Generate a complete set of test certificates for mTLS testing.
///
/// Creates:
/// - CA certificate and key (self-signed root CA)
/// - Server certificate and key (signed by CA, for localhost)
/// - Client certificate and key (signed by CA, for client authentication)
pub fn generate() -> Result<Self, Box<dyn std::error::Error>> {
let temp_dir = TempDir::new()?;
let base_path = temp_dir.path();
// Generate CA key pair
let ca_key = generate_rsa_key()?;
let ca_cert = generate_ca_certificate(&ca_key)?;
// Generate server key pair and certificate
let server_key = generate_rsa_key()?;
let server_cert = generate_server_certificate(&server_key, &ca_cert, &ca_key)?;
// Generate client key pair and certificate
let client_key = generate_rsa_key()?;
let client_cert = generate_client_certificate(&client_key, &ca_cert, &ca_key)?;
// Write all files
let ca_cert_path = base_path.join("ca_cert.pem");
let ca_key_path = base_path.join("ca_key.pem");
let server_cert_path = base_path.join("server_cert.pem");
let server_key_path = base_path.join("server_key.pem");
let client_cert_path = base_path.join("client_cert.pem");
let client_key_path = base_path.join("client_key.pem");
std::fs::write(&ca_cert_path, ca_cert.to_pem()?)?;
std::fs::write(&ca_key_path, ca_key.private_key_to_pem_pkcs8()?)?;
std::fs::write(&server_cert_path, server_cert.to_pem()?)?;
std::fs::write(&server_key_path, server_key.private_key_to_pem_pkcs8()?)?;
std::fs::write(&client_cert_path, client_cert.to_pem()?)?;
std::fs::write(&client_key_path, client_key.private_key_to_pem_pkcs8()?)?;
Ok(Self {
temp_dir,
ca_cert_path,
ca_key_path,
server_cert_path,
server_key_path,
client_cert_path,
client_key_path,
})
}
/// Get paths as string references for use with RouterConfig builder
pub fn ca_cert_str(&self) -> &str {
self.ca_cert_path.to_str().unwrap()
}
pub fn server_cert_str(&self) -> &str {
self.server_cert_path.to_str().unwrap()
}
pub fn server_key_str(&self) -> &str {
self.server_key_path.to_str().unwrap()
}
pub fn client_cert_str(&self) -> &str {
self.client_cert_path.to_str().unwrap()
}
pub fn client_key_str(&self) -> &str {
self.client_key_path.to_str().unwrap()
}
}
/// Generate a 2048-bit RSA key pair
fn generate_rsa_key() -> Result<PKey<Private>, Box<dyn std::error::Error>> {
let rsa = Rsa::generate(2048)?;
Ok(PKey::from_rsa(rsa)?)
}
/// Generate a self-signed CA certificate
fn generate_ca_certificate(key: &PKey<Private>) -> Result<X509, Box<dyn std::error::Error>> {
let mut name_builder = X509NameBuilder::new()?;
name_builder.append_entry_by_text("C", "US")?;
name_builder.append_entry_by_text("ST", "California")?;
name_builder.append_entry_by_text("L", "Test City")?;
name_builder.append_entry_by_text("O", "Test CA Organization")?;
name_builder.append_entry_by_text("CN", "Test CA")?;
let name = name_builder.build();
let mut cert_builder = X509Builder::new()?;
cert_builder.set_version(2)?; // X509 v3
// Serial number
let serial = {
let mut bn = BigNum::new()?;
bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
bn.to_asn1_integer()?
};
cert_builder.set_serial_number(&serial)?;
cert_builder.set_subject_name(&name)?;
cert_builder.set_issuer_name(&name)?; // Self-signed
cert_builder.set_pubkey(key)?;
// Validity: 1 year from now
let not_before = Asn1Time::days_from_now(0)?;
let not_after = Asn1Time::days_from_now(365)?;
cert_builder.set_not_before(&not_before)?;
cert_builder.set_not_after(&not_after)?;
// Extensions for CA
let basic_constraints = BasicConstraints::new().critical().ca().build()?;
cert_builder.append_extension(basic_constraints)?;
let key_usage = KeyUsage::new()
.critical()
.key_cert_sign()
.crl_sign()
.build()?;
cert_builder.append_extension(key_usage)?;
cert_builder.sign(key, MessageDigest::sha256())?;
Ok(cert_builder.build())
}
/// Generate a server certificate signed by the CA
fn generate_server_certificate(
key: &PKey<Private>,
ca_cert: &X509,
ca_key: &PKey<Private>,
) -> Result<X509, Box<dyn std::error::Error>> {
let mut name_builder = X509NameBuilder::new()?;
name_builder.append_entry_by_text("C", "US")?;
name_builder.append_entry_by_text("ST", "California")?;
name_builder.append_entry_by_text("L", "Test City")?;
name_builder.append_entry_by_text("O", "Test Server Organization")?;
name_builder.append_entry_by_text("CN", "localhost")?;
let name = name_builder.build();
let mut cert_builder = X509Builder::new()?;
cert_builder.set_version(2)?;
let serial = {
let mut bn = BigNum::new()?;
bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
bn.to_asn1_integer()?
};
cert_builder.set_serial_number(&serial)?;
cert_builder.set_subject_name(&name)?;
cert_builder.set_issuer_name(ca_cert.subject_name())?;
cert_builder.set_pubkey(key)?;
let not_before = Asn1Time::days_from_now(0)?;
let not_after = Asn1Time::days_from_now(365)?;
cert_builder.set_not_before(&not_before)?;
cert_builder.set_not_after(&not_after)?;
// Extensions for server certificate
let basic_constraints = BasicConstraints::new().build()?;
cert_builder.append_extension(basic_constraints)?;
let key_usage = KeyUsage::new()
.critical()
.digital_signature()
.key_encipherment()
.build()?;
cert_builder.append_extension(key_usage)?;
let ext_key_usage = ExtendedKeyUsage::new().server_auth().build()?;
cert_builder.append_extension(ext_key_usage)?;
// Subject Alternative Names for localhost
let san = SubjectAlternativeName::new()
.dns("localhost")
.ip("127.0.0.1")
.ip("::1")
.build(&cert_builder.x509v3_context(Some(ca_cert), None))?;
cert_builder.append_extension(san)?;
cert_builder.sign(ca_key, MessageDigest::sha256())?;
Ok(cert_builder.build())
}
/// Generate a client certificate signed by the CA
fn generate_client_certificate(
key: &PKey<Private>,
ca_cert: &X509,
ca_key: &PKey<Private>,
) -> Result<X509, Box<dyn std::error::Error>> {
let mut name_builder = X509NameBuilder::new()?;
name_builder.append_entry_by_text("C", "US")?;
name_builder.append_entry_by_text("ST", "California")?;
name_builder.append_entry_by_text("L", "Test City")?;
name_builder.append_entry_by_text("O", "Test Client Organization")?;
name_builder.append_entry_by_text("CN", "Test Client")?;
let name = name_builder.build();
let mut cert_builder = X509Builder::new()?;
cert_builder.set_version(2)?;
let serial = {
let mut bn = BigNum::new()?;
bn.rand(128, MsbOption::MAYBE_ZERO, false)?;
bn.to_asn1_integer()?
};
cert_builder.set_serial_number(&serial)?;
cert_builder.set_subject_name(&name)?;
cert_builder.set_issuer_name(ca_cert.subject_name())?;
cert_builder.set_pubkey(key)?;
let not_before = Asn1Time::days_from_now(0)?;
let not_after = Asn1Time::days_from_now(365)?;
cert_builder.set_not_before(&not_before)?;
cert_builder.set_not_after(&not_after)?;
// Extensions for client certificate
let basic_constraints = BasicConstraints::new().build()?;
cert_builder.append_extension(basic_constraints)?;
let key_usage = KeyUsage::new().critical().digital_signature().build()?;
cert_builder.append_extension(key_usage)?;
let ext_key_usage = ExtendedKeyUsage::new().client_auth().build()?;
cert_builder.append_extension(ext_key_usage)?;
cert_builder.sign(ca_key, MessageDigest::sha256())?;
Ok(cert_builder.build())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_certificate_generation() {
let certs = TestCertificates::generate().expect("Failed to generate certificates");
// Verify all files exist
assert!(certs.ca_cert_path.exists(), "CA cert should exist");
assert!(certs.ca_key_path.exists(), "CA key should exist");
assert!(certs.server_cert_path.exists(), "Server cert should exist");
assert!(certs.server_key_path.exists(), "Server key should exist");
assert!(certs.client_cert_path.exists(), "Client cert should exist");
assert!(certs.client_key_path.exists(), "Client key should exist");
// Verify files are not empty
assert!(
std::fs::metadata(&certs.ca_cert_path).unwrap().len() > 0,
"CA cert should not be empty"
);
assert!(
std::fs::metadata(&certs.server_cert_path).unwrap().len() > 0,
"Server cert should not be empty"
);
assert!(
std::fs::metadata(&certs.client_cert_path).unwrap().len() > 0,
"Client cert should not be empty"
);
}
}
@@ -0,0 +1,316 @@
//! Test configuration builders to reduce duplication across tests
//!
//! Provides pre-configured RouterConfig and MockWorkerConfig builders
//! for common test scenarios.
use smg::config::{CircuitBreakerConfig, PolicyConfig, RetryConfig, RouterConfig};
use super::mock_worker::{HealthStatus, MockWorkerConfig, WorkerType};
/// Default test configuration values
pub mod defaults {
pub const HOST: &str = "127.0.0.1";
pub const MAX_PAYLOAD_SIZE: usize = 256 * 1024 * 1024; // 256MB
pub const REQUEST_TIMEOUT_SECS: u64 = 600;
pub const WORKER_STARTUP_TIMEOUT_SECS: u64 = 5;
pub const WORKER_STARTUP_CHECK_INTERVAL_SECS: u64 = 1;
pub const MAX_CONCURRENT_REQUESTS: i32 = 64;
pub const QUEUE_TIMEOUT_SECS: u64 = 60;
}
/// Builder for common test RouterConfig patterns
pub struct TestRouterConfig;
impl TestRouterConfig {
/// Create a basic round-robin config for routing tests
pub fn round_robin(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a random load balancing config
pub fn random(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.random_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a cache-aware config for routing tests
pub fn cache_aware(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.cache_aware_policy(
0.5, // cache_threshold
32, // balance_abs_threshold
1.5, // balance_rel_threshold
60, // eviction_interval_secs
1000, // max_tree_size
)
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a power-of-two config
pub fn power_of_two(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.power_of_two_policy(5) // load_check_interval_secs
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a manual routing config (for sticky routing tests)
pub fn manual(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.policy(PolicyConfig::Manual {
eviction_interval_secs: 60,
max_idle_secs: 3600,
})
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a config with custom concurrent request limit (for rate limiting tests)
pub fn with_concurrency(port: u16, max_concurrent: i32) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(max_concurrent)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a config with custom payload size limit
pub fn with_payload_limit(port: u16, max_payload_size: usize) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(max_payload_size)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.build_unchecked()
}
/// Create a config with short timeouts (for timeout/retry tests)
pub fn with_short_timeouts(port: u16) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(5)
.worker_startup_timeout_secs(2)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(5)
.build_unchecked()
}
/// Create a round-robin config with retry settings
pub fn round_robin_with_retry(port: u16, retry_config: RetryConfig) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.retry_config(retry_config)
.build_unchecked()
}
/// Create a round-robin config with circuit breaker
pub fn round_robin_with_circuit_breaker(
port: u16,
circuit_breaker: CircuitBreakerConfig,
) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.circuit_breaker_config(circuit_breaker)
.build_unchecked()
}
/// Create a round-robin config with both retry and circuit breaker
pub fn round_robin_with_reliability(
port: u16,
retry_config: RetryConfig,
circuit_breaker: CircuitBreakerConfig,
) -> RouterConfig {
RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host(defaults::HOST)
.port(port)
.max_payload_size(defaults::MAX_PAYLOAD_SIZE)
.request_timeout_secs(defaults::REQUEST_TIMEOUT_SECS)
.worker_startup_timeout_secs(defaults::WORKER_STARTUP_TIMEOUT_SECS)
.worker_startup_check_interval_secs(defaults::WORKER_STARTUP_CHECK_INTERVAL_SECS)
.max_concurrent_requests(defaults::MAX_CONCURRENT_REQUESTS)
.queue_timeout_secs(defaults::QUEUE_TIMEOUT_SECS)
.retry_config(retry_config)
.circuit_breaker_config(circuit_breaker)
.build_unchecked()
}
}
/// Builder for common MockWorkerConfig patterns
pub struct TestWorkerConfig;
impl TestWorkerConfig {
/// Create a healthy worker config
pub fn healthy(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
/// Create multiple healthy workers with sequential ports
pub fn healthy_workers(start_port: u16, count: u16) -> Vec<MockWorkerConfig> {
(0..count).map(|i| Self::healthy(start_port + i)).collect()
}
/// Create an unhealthy worker config
pub fn unhealthy(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Unhealthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
/// Create a slow worker config (for timeout tests)
pub fn slow(port: u16, delay_ms: u64) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: delay_ms,
fail_rate: 0.0,
}
}
/// Create a flaky worker config (for retry/fault tolerance tests)
pub fn flaky(port: u16, fail_rate: f32) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate,
}
}
/// Create a decode worker config (for PD routing tests)
pub fn decode(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Decode,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
/// Create a prefill worker config (for PD routing tests)
pub fn prefill(port: u16) -> MockWorkerConfig {
MockWorkerConfig {
port,
worker_type: WorkerType::Prefill,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_round_robin_config() {
let config = TestRouterConfig::round_robin(3000);
assert_eq!(config.port, 3000);
}
#[test]
fn test_healthy_workers() {
let workers = TestWorkerConfig::healthy_workers(8000, 3);
assert_eq!(workers.len(), 3);
assert_eq!(workers[0].port, 8000);
assert_eq!(workers[1].port, 8001);
assert_eq!(workers[2].port, 8002);
}
}
@@ -0,0 +1,370 @@
// TLS-enabled mock worker for mTLS integration tests
#![allow(dead_code)]
use std::{
net::SocketAddr,
path::Path,
sync::{Arc, Once},
time::{SystemTime, UNIX_EPOCH},
};
use axum::{
extract::{Json, State},
http::StatusCode,
response::{IntoResponse, Response},
routing::{get, post},
Router,
};
use axum_server::tls_rustls::RustlsConfig;
use rustls::{server::WebPkiClientVerifier, RootCertStore};
use serde_json::json;
use tokio::sync::RwLock;
// Ensure crypto provider is installed exactly once
static CRYPTO_PROVIDER_INIT: Once = Once::new();
fn ensure_crypto_provider() {
CRYPTO_PROVIDER_INIT.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
/// Configuration for TLS mock worker behavior
#[derive(Clone)]
pub struct TlsMockWorkerConfig {
pub port: u16,
/// Require client certificate (mTLS) or just server TLS
pub require_client_cert: bool,
/// Response delay in milliseconds
pub response_delay_ms: u64,
/// Fail rate (0.0 - 1.0)
pub fail_rate: f32,
}
impl Default for TlsMockWorkerConfig {
fn default() -> Self {
Self {
port: 0,
require_client_cert: true,
response_delay_ms: 0,
fail_rate: 0.0,
}
}
}
/// TLS-enabled mock worker server for mTLS testing
pub struct TlsMockWorker {
config: Arc<RwLock<TlsMockWorkerConfig>>,
shutdown_handle: Option<tokio::task::JoinHandle<()>>,
shutdown_tx: Option<tokio::sync::oneshot::Sender<()>>,
}
impl TlsMockWorker {
pub fn new(config: TlsMockWorkerConfig) -> Self {
Self {
config: Arc::new(RwLock::new(config)),
shutdown_handle: None,
shutdown_tx: None,
}
}
/// Start the TLS mock worker server
///
/// # Arguments
/// * `server_cert_path` - Path to server certificate PEM file
/// * `server_key_path` - Path to server private key PEM file
/// * `ca_cert_path` - Path to CA certificate for client verification (optional for TLS-only mode)
pub async fn start(
&mut self,
server_cert_path: &Path,
server_key_path: &Path,
ca_cert_path: Option<&Path>,
) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
// Ensure crypto provider is installed before using rustls
ensure_crypto_provider();
let config = self.config.clone();
let port = config.read().await.port;
let require_client_cert = config.read().await.require_client_cert;
// If port is 0, find an available port
let port = if port == 0 {
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let port = listener.local_addr()?.port();
drop(listener);
config.write().await.port = port;
port
} else {
port
};
let app = Router::new()
.route("/health", get(health_handler))
.route("/health_generate", get(health_generate_handler))
.route("/get_server_info", get(server_info_handler))
.route("/generate", post(generate_handler))
.route("/v1/chat/completions", post(chat_completions_handler))
.with_state(config);
let (shutdown_tx, mut shutdown_rx) = tokio::sync::oneshot::channel::<()>();
self.shutdown_tx = Some(shutdown_tx);
// Build TLS configuration
let rustls_config = if require_client_cert {
// mTLS: require client certificate
let ca_cert_path = ca_cert_path.ok_or("CA cert path required for mTLS")?;
build_mtls_config(server_cert_path, server_key_path, ca_cert_path).await?
} else {
// TLS only: no client cert required
build_tls_config(server_cert_path, server_key_path).await?
};
let addr = SocketAddr::from(([127, 0, 0, 1], port));
// Spawn the server in a separate task
let handle = tokio::spawn(async move {
let server =
axum_server::bind_rustls(addr, rustls_config).serve(app.into_make_service());
tokio::select! {
result = server => {
if let Err(e) = result {
eprintln!("TLS Server error: {}", e);
}
}
_ = &mut shutdown_rx => {
// Graceful shutdown
}
}
});
self.shutdown_handle = Some(handle);
// Wait for the server to start
tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
let url = format!("https://127.0.0.1:{}", port);
Ok(url)
}
/// Stop the TLS mock worker server
pub async fn stop(&mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(());
}
if let Some(handle) = self.shutdown_handle.take() {
let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await;
}
}
}
impl Drop for TlsMockWorker {
fn drop(&mut self) {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(());
}
}
}
/// Build TLS config for server-only TLS (no client cert required)
async fn build_tls_config(
cert_path: &Path,
key_path: &Path,
) -> Result<RustlsConfig, Box<dyn std::error::Error + Send + Sync>> {
let config = RustlsConfig::from_pem_file(cert_path, key_path).await?;
Ok(config)
}
/// Build mTLS config requiring client certificate
async fn build_mtls_config(
cert_path: &Path,
key_path: &Path,
ca_cert_path: &Path,
) -> Result<RustlsConfig, Box<dyn std::error::Error + Send + Sync>> {
use std::io::BufReader;
use rustls_pemfile::{certs, pkcs8_private_keys};
// Read server certificate
let cert_file = std::fs::File::open(cert_path)?;
let mut reader = BufReader::new(cert_file);
let cert_chain: Vec<_> = certs(&mut reader).filter_map(|r| r.ok()).collect();
// Read server private key
let key_file = std::fs::File::open(key_path)?;
let mut reader = BufReader::new(key_file);
let private_key = pkcs8_private_keys(&mut reader)
.next()
.ok_or("No private key found")??;
// Read CA certificate for client verification
let ca_file = std::fs::File::open(ca_cert_path)?;
let mut reader = BufReader::new(ca_file);
let ca_certs: Vec<_> = certs(&mut reader).filter_map(|r| r.ok()).collect();
// Build root certificate store for client verification
let mut root_store = RootCertStore::empty();
for cert in ca_certs {
root_store.add(cert)?;
}
// Create client certificate verifier
let client_verifier = WebPkiClientVerifier::builder(Arc::new(root_store))
.build()
.map_err(|e| format!("Failed to build client verifier: {}", e))?;
// Build server config with client verification
let server_config = rustls::ServerConfig::builder()
.with_client_cert_verifier(client_verifier)
.with_single_cert(cert_chain, private_key.into())
.map_err(|e| format!("Failed to build server config: {}", e))?;
Ok(RustlsConfig::from_config(Arc::new(server_config)))
}
// Handler implementations (simplified versions of mock_worker handlers)
async fn should_fail(config: &TlsMockWorkerConfig) -> bool {
rand::random::<f32>() < config.fail_rate
}
async fn health_handler(State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
Json(json!({
"status": "healthy",
"timestamp": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(),
"tls_enabled": true
}))
.into_response()
}
async fn health_generate_handler(
State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
Json(json!({
"status": "ok",
"queue_length": 0,
"processing_time_ms": config.response_delay_ms
}))
.into_response()
}
async fn server_info_handler(State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
Json(json!({
"model_path": "mock-tls-model",
"port": config.port,
"host": "127.0.0.1",
"tls_enabled": true,
"version": "0.3.0"
}))
.into_response()
}
async fn generate_handler(
State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
Json(_payload): Json<serde_json::Value>,
) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({ "error": "Random failure" })),
)
.into_response();
}
if config.response_delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await;
}
Json(json!({
"text": "This is a mock TLS response.",
"meta_info": {
"prompt_tokens": 10,
"completion_tokens": 5,
"tls_verified": true
}
}))
.into_response()
}
async fn chat_completions_handler(
State(config): State<Arc<RwLock<TlsMockWorkerConfig>>>,
Json(_payload): Json<serde_json::Value>,
) -> Response {
let config = config.read().await;
if should_fail(&config).await {
return (
StatusCode::INTERNAL_SERVER_ERROR,
Json(json!({
"error": {
"message": "Random failure",
"type": "internal_error"
}
})),
)
.into_response();
}
if config.response_delay_ms > 0 {
tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await;
}
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
Json(json!({
"id": format!("chatcmpl-{}", uuid::Uuid::new_v4()),
"object": "chat.completion",
"created": timestamp,
"model": "mock-tls-model",
"choices": [{
"index": 0,
"message": {
"role": "assistant",
"content": "This is a mock TLS chat response."
},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15
}
}))
.into_response()
}
@@ -0,0 +1,309 @@
//! Header forwarding integration tests
//!
//! Tests for header propagation through the router to workers.
mod common;
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use common::{
mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
AppTestContext,
};
use serde_json::json;
use smg::config::RouterConfig;
use tower::ServiceExt;
#[cfg(test)]
mod header_forwarding_tests {
use super::*;
/// Test that X-Request-Id header is forwarded
#[tokio::test]
async fn test_request_id_forwarding() {
let ctx = AppTestContext::new(vec![MockWorkerConfig {
port: 19400,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}])
.await;
let app = ctx.create_app().await;
let custom_request_id = "test-request-id-12345";
let payload = json!({
"text": "Test header forwarding",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header("x-request-id", custom_request_id)
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Response should have the same request ID
let response_id = resp.headers().get("x-request-id");
assert!(response_id.is_some(), "Response should have x-request-id");
assert_eq!(
response_id.unwrap().to_str().unwrap(),
custom_request_id,
"Request ID should be preserved"
);
ctx.shutdown().await;
}
/// Test custom request ID headers
#[tokio::test]
async fn test_custom_request_id_headers() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.random_policy()
.host("127.0.0.1")
.port(3500)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.request_id_headers(vec![
"custom-trace-id".to_string(),
"x-correlation-id".to_string(),
])
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 19401,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
let custom_trace_id = "my-custom-trace-123";
let payload = json!({
"text": "Test custom headers",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header("custom-trace-id", custom_trace_id)
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Response should use the custom header as request ID
let response_id = resp.headers().get("x-request-id");
assert!(response_id.is_some());
assert_eq!(
response_id.unwrap().to_str().unwrap(),
custom_trace_id,
"Custom header should be used as request ID"
);
ctx.shutdown().await;
}
/// Test correlation ID header forwarding
#[tokio::test]
async fn test_correlation_id_forwarding() {
let ctx = AppTestContext::new(vec![MockWorkerConfig {
port: 19402,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}])
.await;
let app = ctx.create_app().await;
let correlation_id = "correlation-abc-789";
let payload = json!({
"text": "Test correlation ID",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header("x-correlation-id", correlation_id)
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Response should preserve the correlation ID as request ID
let response_id = resp.headers().get("x-request-id");
assert!(response_id.is_some());
assert_eq!(
response_id.unwrap().to_str().unwrap(),
correlation_id,
"Correlation ID should be preserved"
);
ctx.shutdown().await;
}
/// Test that request ID is generated when not provided
#[tokio::test]
async fn test_auto_generated_request_id() {
let ctx = AppTestContext::new(vec![MockWorkerConfig {
port: 19403,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}])
.await;
let app = ctx.create_app().await;
let payload = json!({
"text": "Test auto-generated ID",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Response should have an auto-generated request ID
let response_id = resp.headers().get("x-request-id");
assert!(
response_id.is_some(),
"Response should have auto-generated x-request-id"
);
let id_value = response_id.unwrap().to_str().unwrap();
assert!(!id_value.is_empty(), "Request ID should not be empty");
// For generate endpoint, ID should have 'gnt-' prefix
assert!(
id_value.starts_with("gnt-"),
"Generate endpoint should have gnt- prefix, got: {}",
id_value
);
ctx.shutdown().await;
}
/// Test chat completions request ID format
#[tokio::test]
async fn test_chat_completions_request_id_format() {
let ctx = AppTestContext::new(vec![MockWorkerConfig {
port: 19404,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}])
.await;
let app = ctx.create_app().await;
let payload = json!({
"model": "test-model",
"messages": [{"role": "user", "content": "Hello"}],
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/v1/chat/completions")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Response should have chatcmpl- prefix for chat completions
let response_id = resp.headers().get("x-request-id");
assert!(response_id.is_some());
let id_value = response_id.unwrap().to_str().unwrap();
assert!(
id_value.starts_with("chatcmpl-"),
"Chat completions should have chatcmpl- prefix, got: {}",
id_value
);
ctx.shutdown().await;
}
/// Test multiple header priorities
#[tokio::test]
async fn test_header_priority() {
let ctx = AppTestContext::new(vec![MockWorkerConfig {
port: 19405,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}])
.await;
let app = ctx.create_app().await;
let primary_id = "primary-request-id";
let fallback_id = "fallback-correlation-id";
let payload = json!({
"text": "Test header priority",
"stream": false
});
// When x-request-id is provided, it should take priority
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header("x-request-id", primary_id)
.header("x-correlation-id", fallback_id)
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
let response_id = resp.headers().get("x-request-id");
assert!(response_id.is_some());
assert_eq!(
response_id.unwrap().to_str().unwrap(),
primary_id,
"x-request-id should take priority"
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,298 @@
//! Payload size integration tests
//!
//! Tests for request payload size limits and handling.
mod common;
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use common::{
mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
AppTestContext,
};
use serde_json::json;
use smg::config::RouterConfig;
use tower::ServiceExt;
#[cfg(test)]
mod payload_size_tests {
use super::*;
/// Test that small payloads are handled correctly
#[tokio::test]
async fn test_small_payload() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4200)
.max_payload_size(1024 * 1024) // 1MB limit
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20200,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
let payload = json!({
"text": "Small payload test",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Small payload should be accepted"
);
ctx.shutdown().await;
}
/// Test that payloads within limit are accepted
#[tokio::test]
async fn test_payload_within_limit() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4201)
.max_payload_size(1024 * 1024) // 1MB limit
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20201,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Create a ~100KB payload (well within 1MB limit)
let large_text = "x".repeat(100 * 1024);
let payload = json!({
"text": large_text,
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Payload within limit should be accepted"
);
ctx.shutdown().await;
}
/// Test that payloads exceeding limit are rejected
#[tokio::test]
async fn test_payload_exceeds_limit() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4202)
.max_payload_size(1024) // Very small 1KB limit
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20202,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Create a payload larger than 1KB limit
let large_text = "x".repeat(2048);
let payload = json!({
"text": large_text,
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
// Should be rejected with 413 Payload Too Large or similar
assert!(
resp.status() == StatusCode::PAYLOAD_TOO_LARGE
|| resp.status() == StatusCode::BAD_REQUEST,
"Payload exceeding limit should be rejected, got {}",
resp.status()
);
ctx.shutdown().await;
}
/// Test edge case: payload exactly at limit
#[tokio::test]
async fn test_payload_at_exact_limit() {
// Use a more reasonable limit for this test
let limit_bytes = 10 * 1024; // 10KB limit
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4203)
.max_payload_size(limit_bytes)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20203,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Create a payload slightly under the limit (accounting for JSON overhead)
let text_size = limit_bytes - 100; // Leave room for JSON structure
let text = "x".repeat(text_size);
let payload = json!({
"text": text,
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
// Payload at/near limit should be accepted
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::PAYLOAD_TOO_LARGE,
"Payload at limit boundary, got status {}",
resp.status()
);
ctx.shutdown().await;
}
/// Test default payload size limit (256MB)
#[tokio::test]
async fn test_default_payload_limit() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4204)
.max_payload_size(256 * 1024 * 1024) // Default 256MB
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20204,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Create a 1MB payload (well within 256MB)
let large_text = "x".repeat(1024 * 1024);
let payload = json!({
"text": large_text,
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"1MB payload should be accepted with 256MB limit"
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,243 @@
//! Circuit breaker integration tests
//!
//! Tests for circuit breaker behavior: opening, half-open state, recovery, and isolation.
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::{CircuitBreakerConfig, RetryConfig, RouterConfig};
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
#[cfg(test)]
mod circuit_breaker_tests {
use super::*;
/// Test that circuit breaker opens after consecutive failures
#[tokio::test]
async fn test_circuit_breaker_opens_after_failures() {
let config = TestRouterConfig::round_robin_with_circuit_breaker(
3200,
CircuitBreakerConfig {
failure_threshold: 3,
success_threshold: 2,
timeout_duration_secs: 2,
window_duration_secs: 10,
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![TestWorkerConfig::flaky(19100, 1.0)], // Always fail
)
.await;
let app = ctx.create_app().await;
// Make requests until we see failures (500) and then circuit breaker opens (503)
let mut saw_500 = false;
let mut saw_503 = false;
for _ in 0..10 {
let payload = json!({
"text": "Test circuit breaker",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
match resp.status() {
StatusCode::INTERNAL_SERVER_ERROR => saw_500 = true,
StatusCode::SERVICE_UNAVAILABLE => saw_503 = true,
_ => {}
}
}
// Should see 500 (worker error) and eventually 503 (circuit breaker open)
assert!(
saw_500 || saw_503,
"Should see either 500 (worker error) or 503 (circuit breaker open)"
);
ctx.shutdown().await;
}
/// Test circuit breaker with disabled flag
#[tokio::test]
async fn test_circuit_breaker_disabled() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(3201)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.disable_circuit_breaker()
.disable_retries()
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![TestWorkerConfig::flaky(19101, 1.0)], // Always fail
)
.await;
let app = ctx.create_app().await;
// With circuit breaker disabled, should always see 500 (never 503)
for _ in 0..5 {
let payload = json!({
"text": "Test disabled CB",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
// With CB disabled, we expect 500 errors (not 503 from CB)
assert_eq!(
resp.status(),
StatusCode::INTERNAL_SERVER_ERROR,
"Should get 500 with CB disabled, not CB-related 503"
);
}
ctx.shutdown().await;
}
/// Test circuit breaker per-worker isolation
#[tokio::test]
async fn test_circuit_breaker_per_worker_isolation() {
let config = TestRouterConfig::round_robin_with_circuit_breaker(
3202,
CircuitBreakerConfig {
failure_threshold: 2,
success_threshold: 1,
timeout_duration_secs: 2,
window_duration_secs: 10,
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19102, 1.0), // Always fail
TestWorkerConfig::healthy(19103), // Always succeed
],
)
.await;
let app = ctx.create_app().await;
// Send enough requests to trigger CB on the failing worker
// The healthy worker should continue to serve requests
let mut success_count = 0;
for _ in 0..20 {
let payload = json!({
"text": "Test isolation",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// After CB opens on failing worker, healthy worker should handle all requests
// We should see at least some successes
assert!(
success_count > 0,
"Should have some successful requests from healthy worker"
);
ctx.shutdown().await;
}
/// Test circuit breaker with retries enabled
#[tokio::test]
async fn test_circuit_breaker_with_retries() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(3203)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.retry_config(RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
})
.circuit_breaker_config(CircuitBreakerConfig {
failure_threshold: 2,
success_threshold: 1,
timeout_duration_secs: 2,
window_duration_secs: 10,
})
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19104, 1.0), // Always fail
TestWorkerConfig::healthy(19105), // Always succeed
],
)
.await;
let app = ctx.create_app().await;
// With retries, requests should succeed by retrying on healthy worker
let payload = json!({
"text": "Test with retries",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed via retry on healthy worker"
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,309 @@
//! Fault tolerance integration tests
//!
//! Tests for system resilience: worker failures, network issues, and recovery.
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::{CircuitBreakerConfig, RetryConfig};
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
#[cfg(test)]
mod fault_tolerance_tests {
use super::*;
/// Test that requests are rerouted when a worker fails
#[tokio::test]
async fn test_worker_failure_reroute() {
let config = TestRouterConfig::round_robin_with_reliability(
4100,
RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 100,
..Default::default()
},
CircuitBreakerConfig {
failure_threshold: 2,
success_threshold: 1,
timeout_duration_secs: 2,
window_duration_secs: 10,
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(20100, 1.0), // Always fails
TestWorkerConfig::healthy(20101), // Always succeeds
],
)
.await;
let app = ctx.create_app().await;
// Requests should succeed via retry to healthy worker
for i in 0..10 {
let payload = json!({
"text": format!("Fault tolerance test {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed via reroute to healthy worker"
);
}
ctx.shutdown().await;
}
/// Test behavior when all workers are temporarily unavailable
#[tokio::test]
async fn test_all_workers_temporarily_failing() {
let config = TestRouterConfig::round_robin_with_retry(
4101,
RetryConfig {
max_retries: 2,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(20102, 1.0), // Always fails
TestWorkerConfig::flaky(20103, 1.0), // Always fails
],
)
.await;
let app = ctx.create_app().await;
let payload = json!({
"text": "Test with all failing workers",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
// Should fail when all workers are failing
assert!(
resp.status() == StatusCode::INTERNAL_SERVER_ERROR
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Request should fail when all workers are failing, got {}",
resp.status()
);
ctx.shutdown().await;
}
/// Test graceful handling of slow workers
#[tokio::test]
async fn test_slow_worker_handling() {
let config = TestRouterConfig::round_robin(4102);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::slow(20104, 500), // Slow worker
TestWorkerConfig::healthy(20105), // Fast worker
],
)
.await;
let app = ctx.create_app().await;
// Send concurrent requests
let mut handles = Vec::new();
let success_count = Arc::new(AtomicUsize::new(0));
for i in 0..10 {
let app_clone = app.clone();
let success_clone = Arc::clone(&success_count);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Slow worker test {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app_clone.oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_clone.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
// All requests should eventually succeed
assert_eq!(
success_count.load(Ordering::SeqCst),
10,
"All requests should succeed despite slow worker"
);
ctx.shutdown().await;
}
/// Test circuit breaker prevents cascading failures
#[tokio::test]
async fn test_circuit_breaker_prevents_cascade() {
let config = TestRouterConfig::round_robin_with_reliability(
4103,
RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
},
CircuitBreakerConfig {
failure_threshold: 2,
success_threshold: 1,
timeout_duration_secs: 5,
window_duration_secs: 10,
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(20106, 1.0), // Failing worker
TestWorkerConfig::healthy(20107), // Healthy worker
],
)
.await;
let app = ctx.create_app().await;
let mut success_count = 0;
// Send many requests - after CB opens, all should route to healthy worker
for i in 0..20 {
let payload = json!({
"text": format!("Circuit breaker cascade test {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// Most requests should succeed after CB opens on failing worker
assert!(
success_count >= 15,
"Most requests should succeed after circuit breaker opens, got {} successes",
success_count
);
ctx.shutdown().await;
}
/// Test recovery after worker comes back online (simulated via healthy worker)
#[tokio::test]
async fn test_system_stability_under_partial_failure() {
let config = TestRouterConfig::round_robin_with_reliability(
4104,
RetryConfig {
max_retries: 2,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
},
CircuitBreakerConfig {
failure_threshold: 3,
success_threshold: 1,
timeout_duration_secs: 2,
window_duration_secs: 10,
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(20108, 0.5), // 50% failure rate
TestWorkerConfig::healthy(20109), // Always succeeds
],
)
.await;
let app = ctx.create_app().await;
let mut success_count = 0;
// System should maintain stability with partial failures
for i in 0..30 {
let payload = json!({
"text": format!("Stability test {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// With retries and one healthy worker, most should succeed
assert!(
success_count >= 25,
"System should maintain stability under partial failure, got {} successes",
success_count
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,6 @@
//! Reliability integration tests
pub mod circuit_breaker_test;
pub mod fault_tolerance_test;
pub mod rate_limiting_test;
pub mod retries_test;
@@ -0,0 +1,263 @@
//! Rate limiting integration tests
//!
//! Tests for rate limiting and concurrency control.
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::RouterConfig;
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
#[cfg(test)]
mod rate_limiting_tests {
use super::*;
/// Test that concurrent requests are handled within limits
#[tokio::test]
async fn test_concurrent_requests_within_limit() {
let config = TestRouterConfig::with_concurrency(3400, 10);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(19300, 50)]).await;
let app = ctx.create_app().await;
// Send 5 concurrent requests (within limit of 10)
let mut handles = Vec::new();
let success_count = Arc::new(AtomicUsize::new(0));
for i in 0..5 {
let app_clone = app.clone();
let success_clone = Arc::clone(&success_count);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Concurrent request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app_clone.oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_clone.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
// Wait for all requests to complete
for handle in handles {
handle.await.unwrap();
}
// All requests should succeed within the limit
assert_eq!(
success_count.load(Ordering::SeqCst),
5,
"All concurrent requests should succeed within limit"
);
ctx.shutdown().await;
}
/// Test rate limit tokens per second
#[tokio::test]
async fn test_rate_limit_tokens() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.random_policy()
.host("127.0.0.1")
.port(3401)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(100)
.rate_limit_tokens_per_second(50) // 50 tokens/sec
.queue_timeout_secs(60)
.build_unchecked();
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19301)]).await;
let app = ctx.create_app().await;
// Send requests and verify they succeed
let mut success_count = 0;
for i in 0..10 {
let payload = json!({
"text": format!("Rate limited request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// Most requests should succeed (some might be rate limited)
assert!(
success_count >= 5,
"At least half of requests should succeed, got {}",
success_count
);
ctx.shutdown().await;
}
/// Test unlimited concurrent requests when set to 0
#[tokio::test]
async fn test_unlimited_concurrent_requests() {
let config = TestRouterConfig::with_concurrency(3402, 0); // Unlimited
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::slow(19302, 10)]).await;
let app = ctx.create_app().await;
// Send many concurrent requests
let mut handles = Vec::new();
let success_count = Arc::new(AtomicUsize::new(0));
for i in 0..20 {
let app_clone = app.clone();
let success_clone = Arc::clone(&success_count);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Unlimited request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app_clone.oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_clone.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
// All requests should succeed with unlimited concurrency
assert_eq!(
success_count.load(Ordering::SeqCst),
20,
"All requests should succeed with unlimited concurrency"
);
ctx.shutdown().await;
}
/// Test queue behavior when requests exceed capacity
#[tokio::test]
async fn test_queue_behavior() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.random_policy()
.host("127.0.0.1")
.port(3403)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(2) // Very low limit
.queue_size(5) // Small queue
.queue_timeout_secs(1) // Short timeout
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![TestWorkerConfig::slow(19303, 200)], // Longer delay to test queuing
)
.await;
let app = ctx.create_app().await;
// Send requests that will exceed capacity
let mut handles = Vec::new();
let success_count = Arc::new(AtomicUsize::new(0));
let error_count = Arc::new(AtomicUsize::new(0));
for i in 0..10 {
let app_clone = app.clone();
let success_clone = Arc::clone(&success_count);
let error_clone = Arc::clone(&error_count);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Queue test {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app_clone.oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_clone.fetch_add(1, Ordering::SeqCst);
} else {
error_clone.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
// Some requests should succeed, some might timeout or be rejected
let successes = success_count.load(Ordering::SeqCst);
let errors = error_count.load(Ordering::SeqCst);
assert!(
successes > 0,
"At least some requests should succeed, got {} successes and {} errors",
successes,
errors
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,214 @@
//! Retry mechanism integration tests
//!
//! Tests for retry behavior: exponential backoff, max retries, and retry-on-failure scenarios.
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::{RetryConfig, RouterConfig};
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
#[cfg(test)]
mod retry_tests {
use super::*;
/// Test that retries succeed when at least one worker is healthy
#[tokio::test]
async fn test_retry_succeeds_with_healthy_fallback() {
let retry_config = RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 100,
..Default::default()
};
let config = TestRouterConfig::round_robin_with_retry(3300, retry_config);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19200, 1.0), // First worker always fails
TestWorkerConfig::healthy(19201), // Second worker always succeeds
],
)
.await;
let app = ctx.create_app().await;
// Request should succeed via retry to healthy worker
let payload = json!({
"text": "Test retry to healthy worker",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed via retry"
);
ctx.shutdown().await;
}
/// Test that retries are disabled when configured
#[tokio::test]
async fn test_retries_disabled() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(3301)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.disable_retries()
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![TestWorkerConfig::flaky(19202, 1.0)], // Always fail
)
.await;
let app = ctx.create_app().await;
// With retries disabled, request should fail immediately
let payload = json!({
"text": "Test no retries",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::INTERNAL_SERVER_ERROR,
"Request should fail without retries"
);
ctx.shutdown().await;
}
/// Test max retries limit
#[tokio::test]
async fn test_max_retries_limit() {
let retry_config = RetryConfig {
max_retries: 2,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
};
let config = TestRouterConfig::round_robin_with_retry(3302, retry_config);
let ctx = AppTestContext::new_with_config(
config,
vec![TestWorkerConfig::flaky(19203, 1.0)], // Always fail
)
.await;
let app = ctx.create_app().await;
// All retries will fail, should return error after exhausting retries
let payload = json!({
"text": "Test max retries",
"stream": false
});
let start = std::time::Instant::now();
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
let elapsed = start.elapsed();
// Should eventually fail after retries
assert!(
resp.status() == StatusCode::INTERNAL_SERVER_ERROR
|| resp.status() == StatusCode::SERVICE_UNAVAILABLE,
"Should fail after exhausting retries, got {}",
resp.status()
);
// Should take some time due to backoff (at least initial_backoff_ms)
// With 2 retries and 10ms initial backoff, should take at least 10ms
// But don't make this too strict as timing can vary
assert!(
elapsed.as_millis() >= 5,
"Should have some backoff delay, got {}ms",
elapsed.as_millis()
);
ctx.shutdown().await;
}
/// Test retry with multiple workers - should eventually find healthy one
#[tokio::test]
async fn test_retry_finds_healthy_worker() {
let retry_config = RetryConfig {
max_retries: 5,
initial_backoff_ms: 5,
max_backoff_ms: 50,
..Default::default()
};
let config = TestRouterConfig::round_robin_with_retry(3303, retry_config);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19204, 1.0), // Fail
TestWorkerConfig::flaky(19205, 1.0), // Fail
TestWorkerConfig::healthy(19206), // Succeed
],
)
.await;
let app = ctx.create_app().await;
// With round robin and retries, should eventually hit the healthy worker
let payload = json!({
"text": "Test find healthy worker",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Should succeed by retrying until finding healthy worker"
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,8 @@
//! Reliability integration tests
#[path = "common/mod.rs"]
pub mod common;
mod reliability;
pub use reliability::*;
@@ -0,0 +1,373 @@
//! Load balancing integration tests
//!
//! Tests for various load balancing policies: round_robin, random, cache_aware, etc.
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::{CircuitBreakerConfig, RetryConfig};
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
#[cfg(test)]
mod round_robin_tests {
use super::*;
/// Test that round robin distributes requests evenly across workers
#[tokio::test]
async fn test_round_robin_distribution() {
let config = TestRouterConfig::round_robin(3100);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19001, 3))
.await;
let app = ctx.create_app().await;
let num_requests = 30;
let mut success_count = 0;
for i in 0..num_requests {
let payload = json!({
"text": format!("Request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// All requests should succeed with round robin across 3 healthy workers
assert_eq!(
success_count, num_requests,
"All requests should succeed with round robin"
);
ctx.shutdown().await;
}
/// Test round robin with one worker failing
#[tokio::test]
async fn test_round_robin_with_failing_worker() {
let config = TestRouterConfig::round_robin_with_retry(
3101,
RetryConfig {
max_retries: 2,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19004, 1.0), // Always fail
TestWorkerConfig::healthy(19005), // Always succeed
],
)
.await;
let app = ctx.create_app().await;
// With retries enabled, requests should eventually succeed
// by being retried on the healthy worker
let payload = json!({
"text": "Test with retries",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed via retry on healthy worker"
);
ctx.shutdown().await;
}
}
#[cfg(test)]
mod random_tests {
use super::*;
/// Test that random policy distributes requests across workers
#[tokio::test]
async fn test_random_distribution() {
let config = TestRouterConfig::random(3102);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19010, 2))
.await;
let app = ctx.create_app().await;
let num_requests = 20;
let mut success_count = 0;
for i in 0..num_requests {
let payload = json!({
"text": format!("Random request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// All requests should succeed with random policy
assert_eq!(
success_count, num_requests,
"All requests should succeed with random policy"
);
ctx.shutdown().await;
}
}
#[cfg(test)]
mod cache_aware_tests {
use super::*;
/// Test cache-aware routing uses consistent hashing
#[tokio::test]
async fn test_cache_aware_consistent_routing() {
let config = TestRouterConfig::cache_aware(3103);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19020, 2))
.await;
let app = ctx.create_app().await;
// Same prompt should route to same worker (consistent hashing)
let same_prompt = "Hello, cache-aware routing test!";
let mut worker_ids: Vec<Option<String>> = Vec::new();
for _ in 0..5 {
let payload = json!({
"text": same_prompt,
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Extract worker ID from response header if available
let worker_id = resp
.headers()
.get("x-worker-id")
.map(|v| v.to_str().unwrap().to_string());
worker_ids.push(worker_id);
}
// All requests should succeed
assert_eq!(worker_ids.len(), 5);
ctx.shutdown().await;
}
/// Test cache-aware routing with different prompts
#[tokio::test]
async fn test_cache_aware_different_prompts() {
let config = TestRouterConfig::cache_aware(3104);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19022, 2))
.await;
let app = ctx.create_app().await;
let prompts = vec![
"First unique prompt",
"Second unique prompt",
"Third unique prompt",
"Fourth unique prompt",
];
let mut success_count = 0;
for prompt in prompts {
let payload = json!({
"text": prompt,
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// All requests should succeed
assert_eq!(success_count, 4);
ctx.shutdown().await;
}
}
#[cfg(test)]
mod worker_health_tests {
use super::*;
/// Test that failing workers are avoided via circuit breaker/retry
#[tokio::test]
async fn test_skip_failing_workers() {
let config = TestRouterConfig::round_robin_with_reliability(
3105,
RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
},
CircuitBreakerConfig {
failure_threshold: 2,
success_threshold: 1,
timeout_duration_secs: 2,
window_duration_secs: 10,
},
);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19030, 1.0), // Always fails
TestWorkerConfig::healthy(19031), // Always succeeds
],
)
.await;
let app = ctx.create_app().await;
// Requests should succeed by retrying on healthy worker
// or by circuit breaker opening on failing worker
for i in 0..10 {
let payload = json!({
"text": format!("Request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
}
ctx.shutdown().await;
}
/// Test behavior when all workers are unhealthy
#[tokio::test]
async fn test_all_workers_unhealthy() {
let ctx = AppTestContext::new(vec![]).await;
let app = ctx.create_app().await;
let payload = json!({
"text": "Test with no workers",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
// Should return service unavailable when no workers available
assert!(
resp.status() == StatusCode::SERVICE_UNAVAILABLE
|| resp.status() == StatusCode::INTERNAL_SERVER_ERROR,
"Expected 503 or 500 when no workers available, got {}",
resp.status()
);
ctx.shutdown().await;
}
}
#[cfg(test)]
mod worker_response_delay_tests {
use super::*;
/// Test that response delays are handled correctly
#[tokio::test]
async fn test_worker_with_delay() {
let config = TestRouterConfig::random(3106);
let ctx = AppTestContext::new_with_config(
config,
vec![TestWorkerConfig::slow(19040, 100)], // 100ms delay
)
.await;
let app = ctx.create_app().await;
let start = std::time::Instant::now();
let payload = json!({
"text": "Test with delay",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
let elapsed = start.elapsed();
assert_eq!(resp.status(), StatusCode::OK);
// Response should take at least 100ms due to configured delay
assert!(
elapsed.as_millis() >= 100,
"Response should be delayed by at least 100ms, got {}ms",
elapsed.as_millis()
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,174 @@
//! Manual routing policy integration tests
//!
//! Tests for the manual routing policy with sticky sessions using X-SMG-Routing-Key header.
use std::collections::{HashMap, HashSet};
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
const ROUTING_KEY_HEADER: &str = "X-SMG-Routing-Key";
#[cfg(test)]
mod manual_routing_tests {
use super::*;
/// Test sticky routing with X-SMG-Routing-Key header
#[tokio::test]
async fn test_manual_routing_with_header() {
let config = TestRouterConfig::manual(3700);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19700, 2))
.await;
let app = ctx.create_app().await;
// Send requests with different routing keys
let mut key_workers: HashMap<String, HashSet<String>> = HashMap::new();
for key_id in 0..5 {
let routing_key = format!("user-{}", key_id);
for _ in 0..4 {
let payload = json!({
"text": format!("Request for {}", routing_key),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header(ROUTING_KEY_HEADER, &routing_key)
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
// Extract worker ID from response header
if let Some(worker_id) = resp.headers().get("x-worker-id") {
let worker = worker_id.to_str().unwrap().to_string();
key_workers
.entry(routing_key.clone())
.or_default()
.insert(worker);
}
}
}
// Verify sticky: each routing key should route to exactly one worker
for (key, workers) in &key_workers {
assert_eq!(
workers.len(),
1,
"Routing key {} should route to exactly one worker, got {:?}",
key,
workers
);
}
ctx.shutdown().await;
}
/// Test random fallback when no routing key header is provided
#[tokio::test]
async fn test_manual_routing_without_header() {
let config = TestRouterConfig::manual(3701);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19702, 2))
.await;
let app = ctx.create_app().await;
let mut success_count = 0;
// Send requests without routing key - should fall back to random selection
for i in 0..20 {
let payload = json!({
"text": format!("Request without key {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
// No ROUTING_KEY_HEADER
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// All requests should succeed via random fallback
assert_eq!(
success_count, 20,
"All requests should succeed with random fallback when no routing key provided"
);
ctx.shutdown().await;
}
/// Test that same routing key consistently routes to same worker
#[tokio::test]
async fn test_manual_routing_consistency() {
let config = TestRouterConfig::manual(3702);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19704, 3))
.await;
let app = ctx.create_app().await;
let routing_key = "consistent-user-123";
let mut seen_workers: Vec<String> = Vec::new();
// Send multiple requests with same routing key
for i in 0..10 {
let payload = json!({
"text": format!("Consistent request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header(ROUTING_KEY_HEADER, routing_key)
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
if let Some(worker_id) = resp.headers().get("x-worker-id") {
seen_workers.push(worker_id.to_str().unwrap().to_string());
}
}
// All requests should go to the same worker
if !seen_workers.is_empty() {
let first_worker = &seen_workers[0];
for worker in &seen_workers {
assert_eq!(
worker, first_worker,
"All requests with same routing key should go to same worker"
);
}
}
ctx.shutdown().await;
}
}
+9
View File
@@ -0,0 +1,9 @@
//! Routing integration tests
pub mod cache_aware_backward_compat_test;
pub mod load_balancing_test;
pub mod manual_routing_test;
pub mod pd_routing_test;
pub mod power_of_two_test;
pub mod test_openai_routing;
pub mod test_pd_routing;
@@ -0,0 +1,222 @@
//! Prefill/Decode (PD) routing integration tests
//!
//! Tests for prefill-decode disaggregation routing mode.
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::RouterConfig;
use tower::ServiceExt;
use crate::common::{
mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
AppTestContext, TestWorkerConfig,
};
#[cfg(test)]
mod pd_routing_tests {
use super::*;
/// Test basic PD mode routing with prefill and decode workers
#[tokio::test]
async fn test_pd_mode_basic_routing() {
let config = RouterConfig::builder()
.prefill_decode_mode(
vec![
("http://127.0.0.1:19800".to_string(), None),
("http://127.0.0.1:19801".to_string(), None),
],
vec![
"http://127.0.0.1:19802".to_string(),
"http://127.0.0.1:19803".to_string(),
],
)
.power_of_two_policy(1)
.host("127.0.0.1")
.port(3800)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
// Note: For PD mode tests, we need to start prefill and decode workers separately
// The test context will need to handle this specially
let ctx = AppTestContext::new_with_config(
config,
vec![
// Prefill workers
TestWorkerConfig::prefill(19800),
TestWorkerConfig::prefill(19801),
// Decode workers
TestWorkerConfig::decode(19802),
TestWorkerConfig::decode(19803),
],
)
.await;
let app = ctx.create_app().await;
// Send requests and verify they succeed
for i in 0..10 {
let payload = json!({
"text": format!("PD mode request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"PD mode request should succeed"
);
}
ctx.shutdown().await;
}
/// Test PD mode with round robin policy
#[tokio::test]
async fn test_pd_mode_round_robin() {
let config = RouterConfig::builder()
.prefill_decode_mode(
vec![("http://127.0.0.1:19810".to_string(), None)],
vec![
"http://127.0.0.1:19811".to_string(),
"http://127.0.0.1:19812".to_string(),
],
)
.round_robin_policy()
.host("127.0.0.1")
.port(3801)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::prefill(19810),
TestWorkerConfig::decode(19811),
TestWorkerConfig::decode(19812),
],
)
.await;
let app = ctx.create_app().await;
let mut success_count = 0;
for i in 0..20 {
let payload = json!({
"text": format!("PD round robin {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
assert_eq!(
success_count, 20,
"All requests should succeed in PD mode with round robin"
);
ctx.shutdown().await;
}
/// Test PD mode handles worker failures gracefully
#[tokio::test]
async fn test_pd_mode_with_failing_decode_worker() {
use smg::config::RetryConfig;
let config = RouterConfig::builder()
.prefill_decode_mode(
vec![("http://127.0.0.1:19820".to_string(), None)],
vec![
"http://127.0.0.1:19821".to_string(),
"http://127.0.0.1:19822".to_string(),
],
)
.round_robin_policy()
.host("127.0.0.1")
.port(3802)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.retry_config(RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
})
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::prefill(19820),
MockWorkerConfig {
port: 19821,
worker_type: WorkerType::Decode,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 1.0, // Failing decode worker
},
TestWorkerConfig::decode(19822), // Healthy decode worker
],
)
.await;
let app = ctx.create_app().await;
// Request should succeed via retry to healthy decode worker
let payload = json!({
"text": "Test with failing decode worker",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed via retry to healthy decode worker"
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,193 @@
//! Power of Two load balancing integration tests
//!
//! Tests for the Power of Two Choices algorithm that selects the less loaded worker.
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use smg::config::RouterConfig;
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
#[cfg(test)]
mod power_of_two_tests {
use super::*;
/// Test that power of two distributes requests across workers
#[tokio::test]
async fn test_power_of_two_distribution() {
let config = TestRouterConfig::power_of_two(3600);
let ctx =
AppTestContext::new_with_config(config, TestWorkerConfig::healthy_workers(19600, 2))
.await;
let app = ctx.create_app().await;
let num_requests = 20;
let mut success_count = 0;
for i in 0..num_requests {
let payload = json!({
"text": format!("Power of two request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
// All requests should succeed
assert_eq!(
success_count, num_requests,
"All requests should succeed with power of two policy"
);
ctx.shutdown().await;
}
/// Test that power of two prefers less loaded workers
#[tokio::test]
async fn test_power_of_two_prefers_less_loaded() {
let config = TestRouterConfig::power_of_two(3601);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::slow(19602, 200), // Slow worker
TestWorkerConfig::healthy(19603), // Fast worker
],
)
.await;
let app = ctx.create_app().await;
// Send concurrent requests to create load imbalance
let mut handles = Vec::new();
let success_count = Arc::new(AtomicUsize::new(0));
for i in 0..30 {
let app_clone = app.clone();
let success_clone = Arc::clone(&success_count);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Concurrent request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app_clone.oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_clone.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
// All requests should succeed
assert!(
success_count.load(Ordering::SeqCst) >= 25,
"Most requests should succeed with power of two"
);
ctx.shutdown().await;
}
/// Test power of two with failing worker uses retry/CB to route to healthy worker
#[tokio::test]
async fn test_power_of_two_with_failing_worker() {
use smg::config::{CircuitBreakerConfig, RetryConfig};
let retry_config = RetryConfig {
max_retries: 3,
initial_backoff_ms: 10,
max_backoff_ms: 50,
..Default::default()
};
let circuit_breaker = CircuitBreakerConfig {
failure_threshold: 2,
success_threshold: 1,
timeout_duration_secs: 2,
window_duration_secs: 10,
};
let config = RouterConfig::builder()
.regular_mode(vec![])
.power_of_two_policy(1)
.host("127.0.0.1")
.port(3602)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.retry_config(retry_config)
.circuit_breaker_config(circuit_breaker)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::flaky(19604, 1.0), // Always fails
TestWorkerConfig::healthy(19605), // Always succeeds
],
)
.await;
let app = ctx.create_app().await;
// Requests should succeed via retry to healthy worker
for i in 0..10 {
let payload = json!({
"text": format!("Request with failing worker {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed via retry on healthy worker"
);
}
ctx.shutdown().await;
}
}
@@ -37,8 +37,7 @@ use tokio::{
};
use tower::ServiceExt;
mod common;
use common::mock_openai_server::MockOpenAIServer;
use crate::common::mock_openai_server::MockOpenAIServer;
/// Helper function to create a minimal chat completion request for testing
fn create_minimal_chat_request() -> ChatCompletionRequest {
@@ -95,9 +94,9 @@ fn create_minimal_completion_request() -> CompletionRequest {
/// Test basic OpenAI router creation and configuration
#[tokio::test]
async fn test_openai_router_creation() {
let ctx = common::test_app::create_test_app_context().await;
let ctx = crate::common::test_app::create_test_app_context().await;
// Register an external worker before creating the router
common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
let router = OpenAIRouter::new(&ctx).await;
assert!(router.is_ok(), "Router creation should succeed");
@@ -110,8 +109,8 @@ async fn test_openai_router_creation() {
/// Test server info endpoint
#[tokio::test]
async fn test_openai_router_server_info() {
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
let req = Request::builder()
@@ -135,8 +134,8 @@ async fn test_openai_router_server_info() {
async fn test_openai_router_models() {
// Use mock server for deterministic models response
let mock_server = MockOpenAIServer::new().await;
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, &mock_server.base_url(), None);
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(&ctx, &mock_server.base_url(), None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
let req = Request::builder()
@@ -208,8 +207,8 @@ async fn test_openai_router_responses_with_mock() {
let base_url = format!("http://{}", addr);
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-4o-mini"]));
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-4o-mini"]));
let router = OpenAIRouter::new(&ctx).await.unwrap();
// Get storage from context (router uses this, not a separate storage)
@@ -473,8 +472,8 @@ async fn test_openai_router_responses_streaming_with_mock() {
let base_url = format!("http://{}", addr);
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-5-nano"]));
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(&ctx, &base_url, Some(vec!["gpt-5-nano"]));
let router = OpenAIRouter::new(&ctx).await.unwrap();
// Get storage from context and seed a previous response
@@ -581,7 +580,7 @@ async fn test_router_factory_openai_mode() {
let router_config = RouterConfig::new(routing_mode, smg::config::PolicyConfig::Random);
let app_context = common::create_test_context(router_config).await;
let app_context = crate::common::create_test_context(router_config).await;
let router = smg::routers::RouterFactory::create_router(&app_context).await;
assert!(
@@ -596,8 +595,8 @@ async fn test_router_factory_openai_mode() {
/// Test that unsupported endpoints return proper error codes
#[tokio::test]
async fn test_unsupported_endpoints() {
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(&ctx, "https://api.openai.com", None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
let generate_request = GenerateRequest {
@@ -655,9 +654,9 @@ async fn test_openai_router_chat_completion_with_mock() {
let mock_server = MockOpenAIServer::new().await;
let base_url = mock_server.base_url();
let ctx = common::test_app::create_test_app_context().await;
let ctx = crate::common::test_app::create_test_app_context().await;
// Register the mock server worker and create router
common::test_app::register_external_worker(&ctx, &base_url, None);
crate::common::test_app::register_external_worker(&ctx, &base_url, None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
// Create a minimal chat completion request
@@ -691,9 +690,9 @@ async fn test_openai_e2e_with_server() {
let mock_server = MockOpenAIServer::new().await;
let base_url = mock_server.base_url();
let ctx = common::test_app::create_test_app_context().await;
let ctx = crate::common::test_app::create_test_app_context().await;
// Register the mock server worker and create router
common::test_app::register_external_worker(&ctx, &base_url, None);
crate::common::test_app::register_external_worker(&ctx, &base_url, None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
// Create Axum app with chat completions endpoint
@@ -757,8 +756,8 @@ async fn test_openai_e2e_with_server() {
async fn test_openai_router_chat_streaming_with_mock() {
let mock_server = MockOpenAIServer::new().await;
let base_url = mock_server.base_url();
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, &base_url, None);
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(&ctx, &base_url, None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
// Build a streaming chat request
@@ -797,8 +796,12 @@ async fn test_openai_router_chat_streaming_with_mock() {
/// Test circuit breaker functionality
#[tokio::test]
async fn test_openai_router_circuit_breaker() {
let ctx = common::test_app::create_test_app_context().await;
common::test_app::register_external_worker(&ctx, "http://invalid-url-that-will-fail", None);
let ctx = crate::common::test_app::create_test_app_context().await;
crate::common::test_app::register_external_worker(
&ctx,
"http://invalid-url-that-will-fail",
None,
);
let router = OpenAIRouter::new(&ctx).await.unwrap();
let chat_request = create_minimal_chat_request();
@@ -821,9 +824,9 @@ async fn test_openai_router_circuit_breaker() {
/// requests to workers, not for the models endpoint.
#[tokio::test]
async fn test_openai_router_models_from_registry() {
let ctx = common::test_app::create_test_app_context().await;
let ctx = crate::common::test_app::create_test_app_context().await;
// Register a worker with the default model
common::test_app::register_external_worker(&ctx, "https://api.example.com", None);
crate::common::test_app::register_external_worker(&ctx, "https://api.example.com", None);
let router = OpenAIRouter::new(&ctx).await.unwrap();
// Get models - should return the registered model
@@ -1,5 +1,5 @@
#[cfg(test)]
mod test_pd_routing {
mod pd_routing_unit_tests {
use serde_json::json;
use smg::{
app_context::AppContext,
+8
View File
@@ -0,0 +1,8 @@
//! Routing integration tests
#[path = "common/mod.rs"]
pub mod common;
mod routing;
pub use routing::*;
@@ -0,0 +1,243 @@
//! Authentication and authorization integration tests
//!
//! Tests for API key enforcement and access control.
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use serde_json::json;
use tower::ServiceExt;
use crate::common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
const AUTH_HEADER: &str = "Authorization";
#[cfg(test)]
mod auth_tests {
use super::*;
/// Test request without API key when auth is not required
#[tokio::test]
async fn test_no_auth_required() {
let config = TestRouterConfig::round_robin(4300);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20300)]).await;
let app = ctx.create_app().await;
// Request without auth header should succeed when no auth required
let payload = json!({
"text": "Test without auth",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request without auth should succeed when no auth required"
);
ctx.shutdown().await;
}
/// Test request with valid API key format
#[tokio::test]
async fn test_with_api_key_header() {
let config = TestRouterConfig::round_robin(4301);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20301)]).await;
let app = ctx.create_app().await;
// Request with Bearer token header
let payload = json!({
"text": "Test with auth header",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header(AUTH_HEADER, "Bearer test-api-key-12345")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
// Without auth enforcement, request should succeed
assert_eq!(
resp.status(),
StatusCode::OK,
"Request with auth header should succeed"
);
ctx.shutdown().await;
}
/// Test health endpoint doesn't require authentication
#[tokio::test]
async fn test_health_endpoint_no_auth() {
let config = TestRouterConfig::round_robin(4302);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20302)]).await;
let app = ctx.create_app().await;
// Health endpoint should be accessible without auth
let req = Request::builder()
.method("GET")
.uri("/health")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Health endpoint should not require auth"
);
ctx.shutdown().await;
}
/// Test OpenAI-compatible API key header (X-API-Key)
#[tokio::test]
async fn test_openai_api_key_header() {
let config = TestRouterConfig::round_robin(4303);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20303)]).await;
let app = ctx.create_app().await;
// Request with X-API-Key header (OpenAI style)
let payload = json!({
"text": "Test with X-API-Key",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header("X-API-Key", "sk-test-key-12345")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request with X-API-Key should succeed"
);
ctx.shutdown().await;
}
/// Test multiple concurrent authenticated requests
#[tokio::test]
async fn test_concurrent_authenticated_requests() {
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
let config = TestRouterConfig::round_robin(4304);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20304)]).await;
let app = ctx.create_app().await;
let success_count = Arc::new(AtomicUsize::new(0));
let mut handles = Vec::new();
for i in 0..20 {
let app_clone = app.clone();
let success_clone = Arc::clone(&success_count);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Concurrent auth test {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.header(AUTH_HEADER, format!("Bearer test-key-{}", i))
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app_clone.oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_clone.fetch_add(1, Ordering::SeqCst);
}
});
handles.push(handle);
}
for handle in handles {
handle.await.unwrap();
}
assert_eq!(
success_count.load(Ordering::SeqCst),
20,
"All concurrent authenticated requests should succeed"
);
ctx.shutdown().await;
}
}
#[cfg(test)]
mod mtls_tests {
use super::*;
/// Test that TLS configuration options exist
/// Note: Actual mTLS testing would require certificate setup
#[tokio::test]
async fn test_tls_config_available() {
// This test verifies the config builder accepts TLS-related options
// Actual mTLS testing requires certificate infrastructure
let config = TestRouterConfig::round_robin(4305);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(20305)]).await;
let app = ctx.create_app().await;
// Basic request should work (no TLS in test mode)
let payload = json!({
"text": "TLS config test",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(resp.status(), StatusCode::OK);
ctx.shutdown().await;
}
}
+5
View File
@@ -0,0 +1,5 @@
//! Security integration tests
pub mod auth_integration_test;
pub mod auth_test;
pub mod mtls_test;
@@ -0,0 +1,506 @@
//! mTLS (Mutual TLS) integration tests
//!
//! Tests for TLS and mTLS communication between router and workers.
//! Covers:
//! - Successful mTLS communication with client certificates
//! - TLS failure without client certificate when required
//! - TLS-only mode (server authentication only)
//! - TLS failure without CA certificate
use std::{io::BufReader, time::Duration};
use rustls::{ClientConfig, RootCertStore};
use rustls_pemfile::{certs, pkcs8_private_keys};
use serde_json::json;
use crate::common::{
test_certs::TestCertificates,
tls_mock_worker::{TlsMockWorker, TlsMockWorkerConfig},
};
// Ensure crypto provider is installed
fn ensure_crypto_provider() {
use std::sync::Once;
static INIT: Once = Once::new();
INIT.call_once(|| {
let _ = rustls::crypto::ring::default_provider().install_default();
});
}
/// Helper to create a TLS-enabled reqwest client with client certificates
fn create_mtls_client(
test_certs: &TestCertificates,
) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
ensure_crypto_provider();
// Read CA certificate
let ca_file = std::fs::File::open(&test_certs.ca_cert_path)?;
let mut ca_reader = BufReader::new(ca_file);
let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
let mut root_store = RootCertStore::empty();
for cert in ca_certs {
root_store.add(cert)?;
}
// Read client certificate
let cert_file = std::fs::File::open(&test_certs.client_cert_path)?;
let mut cert_reader = BufReader::new(cert_file);
let client_certs: Vec<_> = certs(&mut cert_reader).filter_map(|r| r.ok()).collect();
// Read client key
let key_file = std::fs::File::open(&test_certs.client_key_path)?;
let mut key_reader = BufReader::new(key_file);
let client_key = pkcs8_private_keys(&mut key_reader)
.next()
.ok_or("No private key found")??;
// Build client config with client certificate
let client_config = ClientConfig::builder()
.with_root_certificates(root_store)
.with_client_auth_cert(client_certs, client_key.into())
.map_err(|e| format!("Failed to build client config: {}", e))?;
let client = reqwest::Client::builder()
.use_preconfigured_tls(client_config)
.timeout(Duration::from_secs(10))
.build()?;
Ok(client)
}
/// Helper to create a TLS client without client certificates (for server-auth-only)
fn create_tls_client(
test_certs: &TestCertificates,
) -> Result<reqwest::Client, Box<dyn std::error::Error>> {
ensure_crypto_provider();
// Read CA certificate
let ca_file = std::fs::File::open(&test_certs.ca_cert_path)?;
let mut ca_reader = BufReader::new(ca_file);
let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
let mut root_store = RootCertStore::empty();
for cert in ca_certs {
root_store.add(cert)?;
}
// Build client config without client certificate
let client_config = ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
let client = reqwest::Client::builder()
.use_preconfigured_tls(client_config)
.timeout(Duration::from_secs(10))
.build()?;
Ok(client)
}
/// Helper to create a client without CA certificate (for failure test)
fn create_client_without_ca() -> Result<reqwest::Client, Box<dyn std::error::Error>> {
ensure_crypto_provider();
// Build client with empty root store (will fail to verify server)
let root_store = RootCertStore::empty();
let client_config = ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
let client = reqwest::Client::builder()
.use_preconfigured_tls(client_config)
.timeout(Duration::from_secs(10))
.build()?;
Ok(client)
}
#[cfg(test)]
mod mtls_tests {
use super::*;
/// Test successful mTLS communication between client and TLS worker
///
/// This test verifies that:
/// 1. TLS mock worker starts with mTLS configuration
/// 2. Client with proper certificates can connect and communicate
/// 3. Requests succeed with proper authentication
#[tokio::test]
async fn test_mtls_successful_communication() {
// Generate test certificates
let certs = TestCertificates::generate().expect("Failed to generate test certificates");
// Start mTLS-enabled mock worker
let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
port: 0, // Auto-assign port
require_client_cert: true,
response_delay_ms: 0,
fail_rate: 0.0,
});
let worker_url = worker
.start(
&certs.server_cert_path,
&certs.server_key_path,
Some(&certs.ca_cert_path),
)
.await
.expect("Failed to start mTLS worker");
// Create client with full mTLS credentials
let client = create_mtls_client(&certs).expect("Failed to create mTLS client");
// Test health endpoint
let health_resp = client
.get(format!("{}/health", worker_url))
.send()
.await
.expect("Health request failed");
assert!(
health_resp.status().is_success(),
"Health check should succeed with mTLS: {}",
health_resp.status()
);
let health_json: serde_json::Value = health_resp.json().await.unwrap();
assert_eq!(health_json["status"], "healthy");
assert_eq!(health_json["tls_enabled"], true);
// Test generate endpoint
let payload = json!({
"text": "Test mTLS request",
"stream": false
});
let gen_resp = client
.post(format!("{}/generate", worker_url))
.json(&payload)
.send()
.await
.expect("Generate request failed");
assert!(
gen_resp.status().is_success(),
"Generate should succeed with mTLS: {}",
gen_resp.status()
);
let gen_json: serde_json::Value = gen_resp.json().await.unwrap();
assert!(gen_json["text"].as_str().is_some());
assert_eq!(gen_json["meta_info"]["tls_verified"], true);
// Cleanup
worker.stop().await;
}
/// Test that mTLS worker rejects connections without client certificate
///
/// This test verifies that:
/// 1. mTLS worker requires client certificate
/// 2. Connection without client cert fails
#[tokio::test]
async fn test_mtls_failure_without_client_cert() {
// Generate test certificates
let certs = TestCertificates::generate().expect("Failed to generate test certificates");
// Start mTLS-enabled mock worker (requires client cert)
let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
port: 0,
require_client_cert: true,
response_delay_ms: 0,
fail_rate: 0.0,
});
let worker_url = worker
.start(
&certs.server_cert_path,
&certs.server_key_path,
Some(&certs.ca_cert_path),
)
.await
.expect("Failed to start mTLS worker");
// Create client WITHOUT client certificate (only has CA for server verification)
let client = create_tls_client(&certs).expect("Failed to create TLS-only client");
// Attempt to connect - should fail because no client cert provided
let result = client.get(format!("{}/health", worker_url)).send().await;
// Connection should fail or be rejected
assert!(
result.is_err(),
"Connection should fail without client certificate"
);
// Cleanup
worker.stop().await;
}
/// Test TLS-only mode (server authentication only, no client cert required)
///
/// This test verifies that:
/// 1. TLS worker can operate without requiring client certificates
/// 2. Client can connect with just CA certificate for server verification
#[tokio::test]
async fn test_tls_server_auth_only() {
// Generate test certificates
let certs = TestCertificates::generate().expect("Failed to generate test certificates");
// Start TLS-only mock worker (does NOT require client cert)
let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
port: 0,
require_client_cert: false, // TLS-only mode
response_delay_ms: 0,
fail_rate: 0.0,
});
let worker_url = worker
.start(
&certs.server_cert_path,
&certs.server_key_path,
None, // No CA needed for client verification
)
.await
.expect("Failed to start TLS worker");
// Create client with just CA cert for server verification (no client cert)
let client = create_tls_client(&certs).expect("Failed to create TLS client");
// Test health endpoint - should succeed without client cert
let health_resp = client
.get(format!("{}/health", worker_url))
.send()
.await
.expect("Health request failed");
assert!(
health_resp.status().is_success(),
"Health check should succeed in TLS-only mode: {}",
health_resp.status()
);
let health_json: serde_json::Value = health_resp.json().await.unwrap();
assert_eq!(health_json["status"], "healthy");
// Test chat completions endpoint
let payload = json!({
"model": "mock-tls-model",
"messages": [{"role": "user", "content": "Hello TLS"}]
});
let chat_resp = client
.post(format!("{}/v1/chat/completions", worker_url))
.json(&payload)
.send()
.await
.expect("Chat request failed");
assert!(
chat_resp.status().is_success(),
"Chat should succeed in TLS-only mode: {}",
chat_resp.status()
);
// Cleanup
worker.stop().await;
}
/// Test TLS failure when client doesn't have CA certificate
///
/// This test verifies that:
/// 1. Client cannot verify server without proper CA certificate
/// 2. Connection fails due to certificate verification
#[tokio::test]
async fn test_tls_failure_without_ca_cert() {
// Generate test certificates
let certs = TestCertificates::generate().expect("Failed to generate test certificates");
// Start TLS mock worker
let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
port: 0,
require_client_cert: false,
response_delay_ms: 0,
fail_rate: 0.0,
});
let worker_url = worker
.start(&certs.server_cert_path, &certs.server_key_path, None)
.await
.expect("Failed to start TLS worker");
// Create client WITHOUT CA certificate
let client = create_client_without_ca().expect("Failed to create client without CA");
// Attempt to connect - should fail because cannot verify server cert
let result = client.get(format!("{}/health", worker_url)).send().await;
// Connection should fail due to certificate verification
assert!(
result.is_err(),
"Connection should fail without CA certificate for server verification"
);
// Cleanup
worker.stop().await;
}
/// Test multiple concurrent mTLS requests
///
/// This test verifies that mTLS connections work correctly under concurrent load
#[tokio::test]
async fn test_mtls_concurrent_requests() {
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
// Generate test certificates
let certs = TestCertificates::generate().expect("Failed to generate test certificates");
// Start mTLS-enabled mock worker
let mut worker = TlsMockWorker::new(TlsMockWorkerConfig {
port: 0,
require_client_cert: true,
response_delay_ms: 10, // Small delay to simulate work
fail_rate: 0.0,
});
let worker_url = worker
.start(
&certs.server_cert_path,
&certs.server_key_path,
Some(&certs.ca_cert_path),
)
.await
.expect("Failed to start mTLS worker");
// Create mTLS client
let client = Arc::new(create_mtls_client(&certs).expect("Failed to create mTLS client"));
let success_count = Arc::new(AtomicUsize::new(0));
let worker_url = Arc::new(worker_url);
let mut handles = Vec::new();
// Spawn concurrent requests
for i in 0..10 {
let client_clone = Arc::clone(&client);
let success_clone = Arc::clone(&success_count);
let url_clone = Arc::clone(&worker_url);
let handle = tokio::spawn(async move {
let payload = json!({
"text": format!("Concurrent mTLS request {}", i),
"stream": false
});
let resp = client_clone
.post(format!("{}/generate", url_clone))
.json(&payload)
.send()
.await;
if let Ok(response) = resp {
if response.status().is_success() {
success_clone.fetch_add(1, Ordering::SeqCst);
}
}
});
handles.push(handle);
}
// Wait for all requests to complete
for handle in handles {
handle.await.unwrap();
}
// All requests should succeed
assert_eq!(
success_count.load(Ordering::SeqCst),
10,
"All concurrent mTLS requests should succeed"
);
// Cleanup
worker.stop().await;
}
}
#[cfg(test)]
mod certificate_tests {
use super::*;
/// Test that certificate generation works correctly
#[test]
fn test_certificate_generation() {
let certs = TestCertificates::generate().expect("Failed to generate certificates");
// Verify all files exist and are not empty
assert!(certs.ca_cert_path.exists());
assert!(certs.ca_key_path.exists());
assert!(certs.server_cert_path.exists());
assert!(certs.server_key_path.exists());
assert!(certs.client_cert_path.exists());
assert!(certs.client_key_path.exists());
// Verify PEM format
let ca_cert_content = std::fs::read_to_string(&certs.ca_cert_path).unwrap();
assert!(ca_cert_content.contains("-----BEGIN CERTIFICATE-----"));
assert!(ca_cert_content.contains("-----END CERTIFICATE-----"));
let server_key_content = std::fs::read_to_string(&certs.server_key_path).unwrap();
assert!(server_key_content.contains("-----BEGIN PRIVATE KEY-----"));
assert!(server_key_content.contains("-----END PRIVATE KEY-----"));
}
/// Test that generated certificates can be parsed by rustls
#[test]
fn test_certificate_parsing() {
ensure_crypto_provider();
let test_certs = TestCertificates::generate().expect("Failed to generate certificates");
// Verify CA certificate can be parsed by rustls
let ca_file = std::fs::File::open(&test_certs.ca_cert_path).unwrap();
let mut ca_reader = BufReader::new(ca_file);
let ca_certs: Vec<_> = certs(&mut ca_reader).filter_map(|r| r.ok()).collect();
assert!(!ca_certs.is_empty(), "CA certificate should be parseable");
// Verify client certificate can be parsed
let cert_file = std::fs::File::open(&test_certs.client_cert_path).unwrap();
let mut cert_reader = BufReader::new(cert_file);
let client_certs: Vec<_> = certs(&mut cert_reader).filter_map(|r| r.ok()).collect();
assert!(
!client_certs.is_empty(),
"Client certificate should be parseable"
);
// Verify client key can be parsed
let key_file = std::fs::File::open(&test_certs.client_key_path).unwrap();
let mut key_reader = BufReader::new(key_file);
let keys: Vec<_> = pkcs8_private_keys(&mut key_reader)
.filter_map(|r| r.ok())
.collect();
assert!(!keys.is_empty(), "Client key should be parseable");
// Verify we can build a complete client config
let mut root_store = RootCertStore::empty();
for cert in ca_certs {
root_store
.add(cert)
.expect("Should be able to add CA cert to store");
}
let key = keys.into_iter().next().unwrap();
let config_result = ClientConfig::builder()
.with_root_certificates(root_store)
.with_client_auth_cert(client_certs, key.into());
assert!(
config_result.is_ok(),
"Should be able to create client config with certs"
);
}
}
@@ -0,0 +1,8 @@
//! Security integration tests
#[path = "common/mod.rs"]
pub mod common;
mod security;
pub use security::*;
@@ -0,0 +1,261 @@
//! Service discovery integration tests
//!
//! Tests for service discovery shim functionality for dynamic worker registration.
mod common;
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use common::{
mock_worker::{HealthStatus, MockWorkerConfig, WorkerType},
AppTestContext,
};
use serde_json::json;
use smg::config::RouterConfig;
use tower::ServiceExt;
#[cfg(test)]
mod service_discovery_tests {
use super::*;
/// Test service discovery endpoint responds correctly
#[tokio::test]
async fn test_service_discovery_endpoint() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4000)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20000,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Check if service discovery endpoint exists
let req = Request::builder()
.method("GET")
.uri("/v1/workers")
.body(Body::empty())
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
// Endpoint might return OK with worker list or 404 if not implemented
assert!(
resp.status() == StatusCode::OK || resp.status() == StatusCode::NOT_FOUND,
"Workers endpoint should respond, got {}",
resp.status()
);
ctx.shutdown().await;
}
/// Test worker registration via discovery shim
#[tokio::test]
async fn test_worker_registration() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4001)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20001,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Register a new worker via discovery endpoint
let register_payload = json!({
"url": "http://127.0.0.1:20002",
"weight": 1.0
});
let req = Request::builder()
.method("POST")
.uri("/register_worker")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(
serde_json::to_string(&register_payload).unwrap(),
))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
// Registration might succeed or endpoint might not exist
assert!(
resp.status() == StatusCode::OK
|| resp.status() == StatusCode::ACCEPTED
|| resp.status() == StatusCode::NOT_FOUND,
"Registration should respond appropriately, got {}",
resp.status()
);
ctx.shutdown().await;
}
/// Test worker deregistration via discovery shim
#[tokio::test]
async fn test_worker_deregistration() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4002)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![
MockWorkerConfig {
port: 20003,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
},
MockWorkerConfig {
port: 20004,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
},
],
)
.await;
let app = ctx.create_app().await;
// Deregister a worker via discovery endpoint
let deregister_payload = json!({
"url": "http://127.0.0.1:20003"
});
let req = Request::builder()
.method("POST")
.uri("/deregister_worker")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(
serde_json::to_string(&deregister_payload).unwrap(),
))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert!(
resp.status() == StatusCode::OK
|| resp.status() == StatusCode::ACCEPTED
|| resp.status() == StatusCode::NOT_FOUND,
"Deregistration should respond appropriately, got {}",
resp.status()
);
// Requests should still work with remaining worker
let payload = json!({
"text": "Test after deregistration",
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Request should succeed with remaining worker"
);
ctx.shutdown().await;
}
/// Test health status reporting for discovery
#[tokio::test]
async fn test_health_status_endpoint() {
let config = RouterConfig::builder()
.regular_mode(vec![])
.round_robin_policy()
.host("127.0.0.1")
.port(4003)
.max_payload_size(256 * 1024 * 1024)
.request_timeout_secs(600)
.worker_startup_timeout_secs(5)
.worker_startup_check_interval_secs(1)
.max_concurrent_requests(64)
.queue_timeout_secs(60)
.build_unchecked();
let ctx = AppTestContext::new_with_config(
config,
vec![MockWorkerConfig {
port: 20005,
worker_type: WorkerType::Regular,
health_status: HealthStatus::Healthy,
response_delay_ms: 0,
fail_rate: 0.0,
}],
)
.await;
let app = ctx.create_app().await;
// Check health endpoint
let req = Request::builder()
.method("GET")
.uri("/health")
.body(Body::empty())
.unwrap();
let resp = app.oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"Health endpoint should return OK when workers are healthy"
);
ctx.shutdown().await;
}
}
@@ -0,0 +1,17 @@
//! Tool parser integration tests
pub mod tool_parser_deepseek;
pub mod tool_parser_edge_cases;
pub mod tool_parser_fallback;
pub mod tool_parser_glm47_moe;
pub mod tool_parser_glm4_moe;
pub mod tool_parser_json;
pub mod tool_parser_kimik2;
pub mod tool_parser_llama;
pub mod tool_parser_minimax_m2;
pub mod tool_parser_mistral;
pub mod tool_parser_mixed_edge_cases;
pub mod tool_parser_partial_json;
pub mod tool_parser_pythonic;
pub mod tool_parser_qwen;
pub mod tool_parser_step3;
@@ -2,8 +2,7 @@
use smg::tool_parser::{DeepSeekParser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_deepseek_complete_parsing() {
@@ -4,8 +4,7 @@
use smg::tool_parser::{JsonParser, MistralParser, PythonicParser, QwenParser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_empty_input() {
@@ -2,8 +2,7 @@
use smg::tool_parser::{Glm4MoeParser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_glm47_complete_parsing() {
@@ -2,8 +2,7 @@
use smg::tool_parser::{Glm4MoeParser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_glm4_complete_parsing() {
@@ -5,8 +5,7 @@
use serde_json::json;
use smg::tool_parser::{JsonParser, ToolParser};
mod common;
use common::{create_test_tools, streaming_helpers::*};
use crate::common::{create_test_tools, streaming_helpers::*};
#[tokio::test]
async fn test_simple_json_tool_call() {
@@ -2,8 +2,7 @@
use smg::tool_parser::{KimiK2Parser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_kimik2_complete_parsing() {
@@ -4,8 +4,7 @@
use smg::tool_parser::{LlamaParser, ToolParser};
mod common;
use common::{create_test_tools, streaming_helpers::*};
use crate::common::{create_test_tools, streaming_helpers::*};
#[tokio::test]
async fn test_llama_python_tag_format() {
@@ -2,8 +2,7 @@
use smg::tool_parser::{MinimaxM2Parser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_minimax_complete_parsing() {
@@ -7,8 +7,7 @@ use smg::tool_parser::{
JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser, ToolParser,
};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_mixed_formats_in_text() {
@@ -5,8 +5,7 @@
use serde_json::json;
use smg::tool_parser::{PythonicParser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_pythonic_single_function() {
@@ -5,8 +5,7 @@
use serde_json::json;
use smg::tool_parser::{QwenParser, ToolParser};
mod common;
use common::{create_test_tools, streaming_helpers::*};
use crate::common::{create_test_tools, streaming_helpers::*};
#[tokio::test]
async fn test_qwen_single_tool() {
@@ -2,8 +2,7 @@
use smg::tool_parser::{Step3Parser, ToolParser};
mod common;
use common::create_test_tools;
use crate::common::create_test_tools;
#[tokio::test]
async fn test_step3_complete_parsing() {
@@ -0,0 +1,8 @@
//! Tool parser integration tests
#[path = "common/mod.rs"]
pub mod common;
mod tool_parser;
pub use tool_parser::*;
@@ -0,0 +1,140 @@
//! Worker management integration tests
//!
//! Tests for dynamic worker add/remove operations via management API.
//! The actual worker management API uses:
//! - POST /workers - create a worker
//! - GET /workers - list workers
//! - DELETE /workers/{worker_id} - remove a worker
mod common;
use axum::{
body::Body,
extract::Request,
http::{header::CONTENT_TYPE, StatusCode},
};
use common::{AppTestContext, TestRouterConfig, TestWorkerConfig};
use serde_json::json;
use tower::ServiceExt;
#[cfg(test)]
mod worker_management_tests {
use super::*;
/// Test listing workers via API
#[tokio::test]
async fn test_list_workers() {
let config = TestRouterConfig::round_robin(3900);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::healthy(19900),
TestWorkerConfig::healthy(19901),
],
)
.await;
let app = ctx.create_app().await;
// List workers via GET /workers
let req = Request::builder()
.method("GET")
.uri("/workers")
.body(Body::empty())
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
assert_eq!(
resp.status(),
StatusCode::OK,
"GET /workers should return OK"
);
ctx.shutdown().await;
}
/// Test that routing continues to work with multiple workers
#[tokio::test]
async fn test_routing_with_multiple_workers() {
let config = TestRouterConfig::round_robin(3901);
let ctx = AppTestContext::new_with_config(
config,
vec![
TestWorkerConfig::healthy(19902),
TestWorkerConfig::healthy(19903),
],
)
.await;
let app = ctx.create_app().await;
let mut success_count = 0;
// Verify routing distributes across workers
for i in 0..20 {
let payload = json!({
"text": format!("Test request {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
assert_eq!(
success_count, 20,
"All requests should succeed with multiple workers"
);
ctx.shutdown().await;
}
/// Test that requests continue to work during worker operations
#[tokio::test]
async fn test_requests_during_worker_changes() {
let config = TestRouterConfig::round_robin(3902);
let ctx =
AppTestContext::new_with_config(config, vec![TestWorkerConfig::healthy(19904)]).await;
let app = ctx.create_app().await;
// Send requests and verify they succeed
let mut success_count = 0;
for i in 0..10 {
let payload = json!({
"text": format!("Request during changes {}", i),
"stream": false
});
let req = Request::builder()
.method("POST")
.uri("/generate")
.header(CONTENT_TYPE, "application/json")
.body(Body::from(serde_json::to_string(&payload).unwrap()))
.unwrap();
let resp = app.clone().oneshot(req).await.unwrap();
if resp.status() == StatusCode::OK {
success_count += 1;
}
}
assert_eq!(
success_count, 10,
"All requests should succeed during normal operation"
);
ctx.shutdown().await;
}
}