use super::types::*;
impl Default for AdvancedJobScheduler {
fn default() -> Self {
Self::new()
}
}
#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
use super::*;
use std::collections::{HashMap, HashSet};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::time::{Duration, Instant};
fn create_test_node(id: &str) -> ClusterNode {
ClusterNode {
id: id.to_string(),
address: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080),
capabilities: NodeCapabilities {
cpu_cores: 8,
memory_gb: 16,
gpu_count: 1,
storage_gb: 1000,
network_bandwidth_mbps: 1000,
supported_tasks: HashSet::from(["training".to_string(), "inference".to_string()]),
},
status: NodeStatus::Available,
last_heartbeat: Instant::now(),
load_metrics: LoadMetrics {
cpu_usage: 0.3,
memory_usage: 0.4,
gpu_usage: 0.2,
network_io: 0.1,
disk_io: 0.1,
active_jobs: 1,
queue_size: 0,
},
job_history: Vec::new(),
}
}
fn create_test_job(id: &str) -> DistributedJob {
DistributedJob {
id: id.to_string(),
name: format!("test_job_{id}"),
job_type: JobType::Training,
priority: JobPriority::Normal,
requirements: ResourceRequirements {
min_cpu_cores: 2,
min_memory_gb: 4,
min_gpu_count: 0,
min_storage_gb: 10,
preferred_node_tags: HashSet::new(),
exclusive_access: false,
},
created_at: Instant::now(),
timeout: Duration::from_secs(3600),
retry_count: 0,
dependencies: Vec::new(),
metadata: HashMap::new(),
}
}
#[test]
fn test_cluster_creation() {
let cluster = DistributedCluster::new(ClusterConfig::default());
assert!(cluster.get_nodes().is_empty());
}
#[test]
fn test_node_registration() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
assert!(cluster.register_node(node.clone()).is_ok());
assert_eq!(cluster.get_nodes().len(), 1);
assert_eq!(cluster.get_nodes()[0].id, "node1");
}
#[test]
fn test_job_submission() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
let job = create_test_job("job1");
cluster
.register_node(node)
.expect("operation should succeed");
let job_id = cluster.submit_job(job).expect("operation should succeed");
assert_eq!(job_id, "job1");
assert!(cluster.get_job_status(&job_id).is_some());
}
#[test]
fn test_job_scheduling() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
let job = create_test_job("job1");
cluster
.register_node(node)
.expect("operation should succeed");
cluster.submit_job(job).expect("operation should succeed");
let status = cluster.get_job_status("job1");
assert!(status.is_some());
}
#[test]
fn test_job_cancellation() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
let job = create_test_job("job1");
cluster
.register_node(node)
.expect("operation should succeed");
cluster.submit_job(job).expect("operation should succeed");
assert!(cluster.cancel_job("job1").is_ok());
let execution = cluster.get_job_execution("job1");
assert!(execution.is_some());
assert_eq!(
execution.expect("operation should succeed").status,
JobStatus::Cancelled
);
}
#[test]
fn test_node_heartbeat() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
cluster
.register_node(node)
.expect("operation should succeed");
let new_metrics = LoadMetrics {
cpu_usage: 0.8,
memory_usage: 0.7,
gpu_usage: 0.5,
network_io: 0.3,
disk_io: 0.2,
active_jobs: 2,
queue_size: 1,
};
assert!(cluster.update_heartbeat("node1", new_metrics).is_ok());
let nodes = cluster.get_nodes();
assert_eq!(nodes[0].load_metrics.cpu_usage, 0.8);
assert_eq!(nodes[0].status, NodeStatus::Busy);
}
#[test]
fn test_cluster_stats() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node1 = create_test_node("node1");
let node2 = create_test_node("node2");
cluster
.register_node(node1)
.expect("operation should succeed");
cluster
.register_node(node2)
.expect("operation should succeed");
let stats = cluster.get_cluster_stats();
assert_eq!(stats.total_nodes, 2);
assert_eq!(stats.available_nodes, 2);
assert_eq!(stats.total_cpu_cores, 16);
assert_eq!(stats.total_memory_gb, 32);
}
#[test]
fn test_job_scheduler() {
let scheduler = JobScheduler::new();
let mut nodes = HashMap::new();
let node1 = create_test_node("node1");
let node2 = create_test_node("node2");
nodes.insert("node1".to_string(), node1);
nodes.insert("node2".to_string(), node2);
let job = create_test_job("job1");
let selected_node = scheduler.find_suitable_node(&job, &nodes);
assert!(selected_node.is_some());
assert!(
["node1", "node2"].contains(&selected_node.expect("operation should succeed").as_str())
);
}
#[test]
fn test_load_balancer() {
let load_balancer = LoadBalancer::new();
let mut nodes = HashMap::new();
let node1 = create_test_node("node1");
nodes.insert("node1".to_string(), node1);
assert!(load_balancer.rebalance(&nodes).is_ok());
}
#[test]
fn test_fault_detector() {
let mut fault_detector = FaultDetector::new();
assert!(fault_detector.handle_failure("node1").is_ok());
assert!(!fault_detector.is_problematic("node1"));
for _ in 0..4 {
fault_detector
.handle_failure("node1")
.expect("operation should succeed");
}
assert!(fault_detector.is_problematic("node1"));
}
#[test]
fn test_node_failure_handling() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
let job = create_test_job("job1");
cluster
.register_node(node)
.expect("operation should succeed");
cluster.submit_job(job).expect("operation should succeed");
assert!(cluster.handle_node_failure("node1").is_ok());
let execution = cluster.get_job_execution("job1");
if let Some(exec) = execution {
println!("Job status: {:?}", exec.status);
}
}
#[test]
fn test_resource_requirements() {
let scheduler = JobScheduler::new();
let mut nodes = HashMap::new();
let node = create_test_node("node1");
nodes.insert("node1".to_string(), node);
let mut job = create_test_job("job1");
job.requirements.min_cpu_cores = 16;
let selected_node = scheduler.find_suitable_node(&job, &nodes);
assert!(selected_node.is_none());
job.requirements.min_cpu_cores = 4;
let selected_node = scheduler.find_suitable_node(&job, &nodes);
assert!(selected_node.is_some());
}
#[test]
fn test_job_priorities() {
let cluster = DistributedCluster::new(ClusterConfig::default());
let node = create_test_node("node1");
cluster
.register_node(node)
.expect("operation should succeed");
let mut job1 = create_test_job("job1");
job1.priority = JobPriority::Low;
let mut job2 = create_test_job("job2");
job2.priority = JobPriority::High;
cluster.submit_job(job1).expect("operation should succeed");
cluster.submit_job(job2).expect("operation should succeed");
let queue = cluster.job_queue.lock().expect("operation should succeed");
if !queue.is_empty() {
assert_eq!(queue[0].priority, JobPriority::High);
}
}
#[test]
fn test_message_passing_system() {
let mps = MessagePassingSystem::new("node1".to_string());
let message = DistributedMessage {
id: "msg1".to_string(),
source: "node1".to_string(),
destination: "node2".to_string(),
message_type: MessageType::JobSubmission,
data: vec![1, 2, 3, 4],
timestamp: Instant::now(),
priority: MessagePriority::Normal,
};
assert!(mps.send_message(message.clone()).is_err());
mps.routing_table
.write()
.expect("operation should succeed")
.insert(
"node2".to_string(),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8081),
);
assert!(mps.send_message(message).is_ok());
}
#[test]
fn test_message_broadcasting() {
let mps = MessagePassingSystem::new("node1".to_string());
mps.routing_table
.write()
.expect("operation should succeed")
.insert(
"node2".to_string(),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8081),
);
mps.routing_table
.write()
.expect("operation should succeed")
.insert(
"node3".to_string(),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8082),
);
let data = vec![5, 6, 7, 8];
assert!(mps.broadcast_message(MessageType::Heartbeat, data).is_ok());
let queue = mps.message_queue.lock().expect("operation should succeed");
assert_eq!(queue.len(), 2);
}
#[test]
fn test_message_handler() {
let handler = MessageHandler::new(|msg: &DistributedMessage| {
Ok(MessageResponse {
message_id: msg.id.clone(),
success: true,
data: vec![],
error: None,
})
});
let message = DistributedMessage {
id: "test_msg".to_string(),
source: "node1".to_string(),
destination: "node2".to_string(),
message_type: MessageType::JobSubmission,
data: vec![],
timestamp: Instant::now(),
priority: MessagePriority::Normal,
};
let response = handler.handle(&message).expect("operation should succeed");
assert!(response.success);
assert_eq!(response.message_id, "test_msg");
}
#[test]
fn test_consensus_manager() {
let mut consensus = ConsensusManager::new(
"node1".to_string(),
vec!["node2".to_string(), "node3".to_string()],
);
let (state, term) = consensus.get_state();
assert_eq!(state, ConsensusState::Follower);
assert_eq!(term, 0);
assert!(consensus.start_election().is_ok());
let (state, term) = consensus.get_state();
assert_eq!(state, ConsensusState::Candidate);
assert_eq!(term, 1);
let vote_request = VoteRequest {
term: 2,
candidate_id: "node2".to_string(),
last_log_index: 0,
last_log_term: 0,
};
let response = consensus.handle_vote_request(vote_request);
assert!(response.vote_granted);
assert_eq!(response.term, 2);
}
#[test]
fn test_consensus_log_entry() {
let mut consensus = ConsensusManager::new("node1".to_string(), vec!["node2".to_string()]);
let entry = LogEntry {
term: 1,
index: 0,
command: "test_command".to_string(),
data: vec![1, 2, 3],
};
assert!(consensus.append_entry(entry.clone()).is_err());
consensus.state = ConsensusState::Leader;
assert!(consensus.append_entry(entry).is_ok());
assert_eq!(consensus.log.len(), 1);
}
#[test]
fn test_data_partitioner() {
let mut partitioner = DataPartitioner::new(PartitioningStrategy::Hash, 4, 2);
let partition1 = partitioner.get_partition("key1");
let partition2 = partitioner.get_partition("key2");
assert!(partition1 < 4);
assert!(partition2 < 4);
assert_eq!(partition1, partitioner.get_partition("key1"));
partitioner.assign_partition(0, "node1".to_string());
partitioner.assign_partition(1, "node2".to_string());
let nodes = partitioner.get_partition_nodes(0);
assert!(nodes.contains(&"node1".to_string()));
}
#[test]
fn test_data_partitioner_rebalancing() {
let mut partitioner = DataPartitioner::new(PartitioningStrategy::Hash, 4, 1);
let nodes = vec!["node1".to_string(), "node2".to_string()];
let result = partitioner.rebalance_partitions(&nodes);
assert_eq!(result.assignments_changed, 4);
assert_eq!(result.partitions_moved.len(), 0);
}
#[test]
fn test_partitioning_strategies() {
let hash_partitioner = DataPartitioner::new(PartitioningStrategy::Hash, 4, 1);
let range_partitioner = DataPartitioner::new(PartitioningStrategy::Range, 4, 1);
let random_partitioner = DataPartitioner::new(PartitioningStrategy::Random, 4, 1);
let key = "test_key";
let hash_partition = hash_partitioner.get_partition(key);
let range_partition = range_partitioner.get_partition(key);
let random_partition = random_partitioner.get_partition(key);
assert!(hash_partition < 4);
assert!(range_partition < 4);
assert!(random_partition < 4);
}
#[test]
fn test_advanced_job_scheduler() {
let mut scheduler = AdvancedJobScheduler::new();
let mut nodes = HashMap::new();
let node1 = create_test_node("node1");
let node2 = create_test_node("node2");
nodes.insert("node1".to_string(), node1);
nodes.insert("node2".to_string(), node2);
let jobs = vec![create_test_job("job1"), create_test_job("job2")];
let decisions = scheduler
.gang_schedule(&jobs, &nodes)
.expect("operation should succeed");
assert_eq!(decisions.len(), jobs.len());
for decision in &decisions {
assert!(nodes.contains_key(&decision.node_id));
assert!(decision.resource_allocation.cpu_cores > 0);
}
}
#[test]
fn test_backfill_scheduling() {
let mut scheduler = AdvancedJobScheduler::new();
let mut nodes = HashMap::new();
let node1 = create_test_node("node1");
nodes.insert("node1".to_string(), node1);
let waiting_jobs = vec![create_test_job("waiting_job")];
let decisions = scheduler
.backfill_schedule(&waiting_jobs, &nodes)
.expect("operation should succeed");
assert_eq!(decisions.len(), 1);
assert_eq!(decisions[0].job_id, "waiting_job");
}
#[test]
fn test_resource_reservation() {
let mut scheduler = AdvancedJobScheduler::new();
let reservation = ResourceReservation {
id: "reservation1".to_string(),
node_id: "node1".to_string(),
start_time: Instant::now(),
duration: Duration::from_secs(3600),
resources: ResourceAllocation {
cpu_cores: 4,
memory_gb: 8,
gpu_count: 1,
storage_gb: 100,
network_bandwidth: 1000,
},
};
assert!(scheduler.reserve_resources(reservation).is_ok());
assert_eq!(scheduler.resource_reservations.len(), 1);
}
#[test]
fn test_checkpoint_manager() {
let mut checkpoint_mgr = CheckpointManager::new(Duration::from_secs(300));
let job_state = JobState {
progress: 0.5,
intermediate_results: HashMap::new(),
runtime_state: vec![1, 2, 3, 4],
};
let checkpoint_id = checkpoint_mgr
.create_checkpoint("job1", job_state.clone())
.expect("operation should succeed");
assert!(!checkpoint_id.is_empty());
let restored_state = checkpoint_mgr
.restore_checkpoint(&checkpoint_id)
.expect("operation should succeed");
assert_eq!(restored_state.progress, 0.5);
assert_eq!(restored_state.runtime_state, vec![1, 2, 3, 4]);
let stats = checkpoint_mgr.get_checkpoint_stats();
assert_eq!(stats.total_checkpoints, 1);
assert!(stats.total_size_bytes > 0);
}
#[test]
fn test_checkpoint_cleanup() {
let mut checkpoint_mgr = CheckpointManager::new(Duration::from_secs(300));
let job_state = JobState {
progress: 1.0,
intermediate_results: HashMap::new(),
runtime_state: vec![],
};
checkpoint_mgr
.create_checkpoint("job1", job_state.clone())
.expect("operation should succeed");
checkpoint_mgr
.create_checkpoint("job2", job_state)
.expect("operation should succeed");
assert_eq!(checkpoint_mgr.checkpoint_storage.len(), 2);
checkpoint_mgr.cleanup_old_checkpoints(Duration::from_secs(0));
assert_eq!(checkpoint_mgr.checkpoint_storage.len(), 0);
}
#[test]
fn test_message_type_conversion() {
assert_eq!(format!("{}", MessageType::JobSubmission), "job_submission");
assert_eq!(format!("{}", MessageType::JobResult), "job_result");
assert_eq!(format!("{}", MessageType::Heartbeat), "heartbeat");
assert_eq!(
format!("{}", MessageType::ResourceUpdate),
"resource_update"
);
assert_eq!(
format!("{}", MessageType::ConsensusRequest),
"consensus_request"
);
assert_eq!(format!("{}", MessageType::DataPartition), "data_partition");
assert_eq!(
format!("{}", MessageType::Custom("test".to_string())),
"test"
);
}
#[test]
fn test_consensus_states() {
let follower = ConsensusState::Follower;
let candidate = ConsensusState::Candidate;
let leader = ConsensusState::Leader;
assert_eq!(follower, ConsensusState::Follower);
assert_eq!(candidate, ConsensusState::Candidate);
assert_eq!(leader, ConsensusState::Leader);
}
#[test]
fn test_scheduling_policies() {
let policies = [
SchedulingPolicy::FIFO,
SchedulingPolicy::ShortestJobFirst,
SchedulingPolicy::GangScheduling,
SchedulingPolicy::Backfill,
SchedulingPolicy::PriorityBased,
];
assert_eq!(policies.len(), 5);
}
#[test]
fn test_message_priorities() {
let low = MessagePriority::Low;
let normal = MessagePriority::Normal;
let high = MessagePriority::High;
let critical = MessagePriority::Critical;
match low {
MessagePriority::Low => {}
_ => panic!(),
}
match normal {
MessagePriority::Normal => {}
_ => panic!(),
}
match high {
MessagePriority::High => {}
_ => panic!(),
}
match critical {
MessagePriority::Critical => {}
_ => panic!(),
}
}
#[test]
fn test_resource_allocation_calculations() {
let allocation = ResourceAllocation {
cpu_cores: 8,
memory_gb: 16,
gpu_count: 2,
storage_gb: 500,
network_bandwidth: 1000,
};
assert_eq!(allocation.cpu_cores, 8);
assert_eq!(allocation.memory_gb, 16);
assert_eq!(allocation.gpu_count, 2);
assert_eq!(allocation.storage_gb, 500);
assert_eq!(allocation.network_bandwidth, 1000);
}
#[test]
fn test_job_state_serialization() {
let mut intermediate_results = HashMap::new();
intermediate_results.insert("result1".to_string(), vec![1, 2, 3]);
intermediate_results.insert("result2".to_string(), vec![4, 5, 6]);
let job_state = JobState {
progress: 0.75,
intermediate_results,
runtime_state: vec![7, 8, 9],
};
assert_eq!(job_state.progress, 0.75);
assert_eq!(job_state.intermediate_results.len(), 2);
assert_eq!(job_state.runtime_state, vec![7, 8, 9]);
assert!(job_state.intermediate_results.contains_key("result1"));
assert!(job_state.intermediate_results.contains_key("result2"));
}
}