1use super::types::*;
6impl Default for AdvancedJobScheduler {
7 fn default() -> Self {
8 Self::new()
9 }
10}
11#[allow(non_snake_case)]
12#[cfg(test)]
13mod tests {
14 use super::*;
15 use std::collections::{HashMap, HashSet};
16 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
17 use std::time::{Duration, Instant};
18 fn create_test_node(id: &str) -> ClusterNode {
19 ClusterNode {
20 id: id.to_string(),
21 address: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080),
22 capabilities: NodeCapabilities {
23 cpu_cores: 8,
24 memory_gb: 16,
25 gpu_count: 1,
26 storage_gb: 1000,
27 network_bandwidth_mbps: 1000,
28 supported_tasks: HashSet::from(["training".to_string(), "inference".to_string()]),
29 },
30 status: NodeStatus::Available,
31 last_heartbeat: Instant::now(),
32 load_metrics: LoadMetrics {
33 cpu_usage: 0.3,
34 memory_usage: 0.4,
35 gpu_usage: 0.2,
36 network_io: 0.1,
37 disk_io: 0.1,
38 active_jobs: 1,
39 queue_size: 0,
40 },
41 job_history: Vec::new(),
42 }
43 }
44 fn create_test_job(id: &str) -> DistributedJob {
45 DistributedJob {
46 id: id.to_string(),
47 name: format!("test_job_{id}"),
48 job_type: JobType::Training,
49 priority: JobPriority::Normal,
50 requirements: ResourceRequirements {
51 min_cpu_cores: 2,
52 min_memory_gb: 4,
53 min_gpu_count: 0,
54 min_storage_gb: 10,
55 preferred_node_tags: HashSet::new(),
56 exclusive_access: false,
57 },
58 created_at: Instant::now(),
59 timeout: Duration::from_secs(3600),
60 retry_count: 0,
61 dependencies: Vec::new(),
62 metadata: HashMap::new(),
63 }
64 }
65 #[test]
66 fn test_cluster_creation() {
67 let cluster = DistributedCluster::new(ClusterConfig::default());
68 assert!(cluster.get_nodes().is_empty());
69 }
70 #[test]
71 fn test_node_registration() {
72 let cluster = DistributedCluster::new(ClusterConfig::default());
73 let node = create_test_node("node1");
74 assert!(cluster.register_node(node.clone()).is_ok());
75 assert_eq!(cluster.get_nodes().len(), 1);
76 assert_eq!(cluster.get_nodes()[0].id, "node1");
77 }
78 #[test]
79 fn test_job_submission() {
80 let cluster = DistributedCluster::new(ClusterConfig::default());
81 let node = create_test_node("node1");
82 let job = create_test_job("job1");
83 cluster
84 .register_node(node)
85 .expect("operation should succeed");
86 let job_id = cluster.submit_job(job).expect("operation should succeed");
87 assert_eq!(job_id, "job1");
88 assert!(cluster.get_job_status(&job_id).is_some());
89 }
90 #[test]
91 fn test_job_scheduling() {
92 let cluster = DistributedCluster::new(ClusterConfig::default());
93 let node = create_test_node("node1");
94 let job = create_test_job("job1");
95 cluster
96 .register_node(node)
97 .expect("operation should succeed");
98 cluster.submit_job(job).expect("operation should succeed");
99 let status = cluster.get_job_status("job1");
100 assert!(status.is_some());
101 }
102 #[test]
103 fn test_job_cancellation() {
104 let cluster = DistributedCluster::new(ClusterConfig::default());
105 let node = create_test_node("node1");
106 let job = create_test_job("job1");
107 cluster
108 .register_node(node)
109 .expect("operation should succeed");
110 cluster.submit_job(job).expect("operation should succeed");
111 assert!(cluster.cancel_job("job1").is_ok());
112 let execution = cluster.get_job_execution("job1");
113 assert!(execution.is_some());
114 assert_eq!(
115 execution.expect("operation should succeed").status,
116 JobStatus::Cancelled
117 );
118 }
119 #[test]
120 fn test_node_heartbeat() {
121 let cluster = DistributedCluster::new(ClusterConfig::default());
122 let node = create_test_node("node1");
123 cluster
124 .register_node(node)
125 .expect("operation should succeed");
126 let new_metrics = LoadMetrics {
127 cpu_usage: 0.8,
128 memory_usage: 0.7,
129 gpu_usage: 0.5,
130 network_io: 0.3,
131 disk_io: 0.2,
132 active_jobs: 2,
133 queue_size: 1,
134 };
135 assert!(cluster.update_heartbeat("node1", new_metrics).is_ok());
136 let nodes = cluster.get_nodes();
137 assert_eq!(nodes[0].load_metrics.cpu_usage, 0.8);
138 assert_eq!(nodes[0].status, NodeStatus::Busy);
139 }
140 #[test]
141 fn test_cluster_stats() {
142 let cluster = DistributedCluster::new(ClusterConfig::default());
143 let node1 = create_test_node("node1");
144 let node2 = create_test_node("node2");
145 cluster
146 .register_node(node1)
147 .expect("operation should succeed");
148 cluster
149 .register_node(node2)
150 .expect("operation should succeed");
151 let stats = cluster.get_cluster_stats();
152 assert_eq!(stats.total_nodes, 2);
153 assert_eq!(stats.available_nodes, 2);
154 assert_eq!(stats.total_cpu_cores, 16);
155 assert_eq!(stats.total_memory_gb, 32);
156 }
157 #[test]
158 fn test_job_scheduler() {
159 let scheduler = JobScheduler::new();
160 let mut nodes = HashMap::new();
161 let node1 = create_test_node("node1");
162 let node2 = create_test_node("node2");
163 nodes.insert("node1".to_string(), node1);
164 nodes.insert("node2".to_string(), node2);
165 let job = create_test_job("job1");
166 let selected_node = scheduler.find_suitable_node(&job, &nodes);
167 assert!(selected_node.is_some());
168 assert!(
169 ["node1", "node2"].contains(&selected_node.expect("operation should succeed").as_str())
170 );
171 }
172 #[test]
173 fn test_load_balancer() {
174 let load_balancer = LoadBalancer::new();
175 let mut nodes = HashMap::new();
176 let node1 = create_test_node("node1");
177 nodes.insert("node1".to_string(), node1);
178 assert!(load_balancer.rebalance(&nodes).is_ok());
179 }
180 #[test]
181 fn test_fault_detector() {
182 let mut fault_detector = FaultDetector::new();
183 assert!(fault_detector.handle_failure("node1").is_ok());
184 assert!(!fault_detector.is_problematic("node1"));
185 for _ in 0..4 {
186 fault_detector
187 .handle_failure("node1")
188 .expect("operation should succeed");
189 }
190 assert!(fault_detector.is_problematic("node1"));
191 }
192 #[test]
193 fn test_node_failure_handling() {
194 let cluster = DistributedCluster::new(ClusterConfig::default());
195 let node = create_test_node("node1");
196 let job = create_test_job("job1");
197 cluster
198 .register_node(node)
199 .expect("operation should succeed");
200 cluster.submit_job(job).expect("operation should succeed");
201 assert!(cluster.handle_node_failure("node1").is_ok());
202 let execution = cluster.get_job_execution("job1");
203 if let Some(exec) = execution {
204 println!("Job status: {:?}", exec.status);
205 }
206 }
207 #[test]
208 fn test_resource_requirements() {
209 let scheduler = JobScheduler::new();
210 let mut nodes = HashMap::new();
211 let node = create_test_node("node1");
212 nodes.insert("node1".to_string(), node);
213 let mut job = create_test_job("job1");
214 job.requirements.min_cpu_cores = 16;
215 let selected_node = scheduler.find_suitable_node(&job, &nodes);
216 assert!(selected_node.is_none());
217 job.requirements.min_cpu_cores = 4;
218 let selected_node = scheduler.find_suitable_node(&job, &nodes);
219 assert!(selected_node.is_some());
220 }
221 #[test]
222 fn test_job_priorities() {
223 let cluster = DistributedCluster::new(ClusterConfig::default());
224 let node = create_test_node("node1");
225 cluster
226 .register_node(node)
227 .expect("operation should succeed");
228 let mut job1 = create_test_job("job1");
229 job1.priority = JobPriority::Low;
230 let mut job2 = create_test_job("job2");
231 job2.priority = JobPriority::High;
232 cluster.submit_job(job1).expect("operation should succeed");
233 cluster.submit_job(job2).expect("operation should succeed");
234 let queue = cluster.job_queue.lock().expect("operation should succeed");
235 if !queue.is_empty() {
236 assert_eq!(queue[0].priority, JobPriority::High);
237 }
238 }
239 #[test]
240 fn test_message_passing_system() {
241 let mps = MessagePassingSystem::new("node1".to_string());
242 let message = DistributedMessage {
243 id: "msg1".to_string(),
244 source: "node1".to_string(),
245 destination: "node2".to_string(),
246 message_type: MessageType::JobSubmission,
247 data: vec![1, 2, 3, 4],
248 timestamp: Instant::now(),
249 priority: MessagePriority::Normal,
250 };
251 assert!(mps.send_message(message.clone()).is_err());
252 mps.routing_table
253 .write()
254 .expect("operation should succeed")
255 .insert(
256 "node2".to_string(),
257 SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8081),
258 );
259 assert!(mps.send_message(message).is_ok());
260 }
261 #[test]
262 fn test_message_broadcasting() {
263 let mps = MessagePassingSystem::new("node1".to_string());
264 mps.routing_table
265 .write()
266 .expect("operation should succeed")
267 .insert(
268 "node2".to_string(),
269 SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8081),
270 );
271 mps.routing_table
272 .write()
273 .expect("operation should succeed")
274 .insert(
275 "node3".to_string(),
276 SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8082),
277 );
278 let data = vec![5, 6, 7, 8];
279 assert!(mps.broadcast_message(MessageType::Heartbeat, data).is_ok());
280 let queue = mps.message_queue.lock().expect("operation should succeed");
281 assert_eq!(queue.len(), 2);
282 }
283 #[test]
284 fn test_message_handler() {
285 let handler = MessageHandler::new(|msg: &DistributedMessage| {
286 Ok(MessageResponse {
287 message_id: msg.id.clone(),
288 success: true,
289 data: vec![],
290 error: None,
291 })
292 });
293 let message = DistributedMessage {
294 id: "test_msg".to_string(),
295 source: "node1".to_string(),
296 destination: "node2".to_string(),
297 message_type: MessageType::JobSubmission,
298 data: vec![],
299 timestamp: Instant::now(),
300 priority: MessagePriority::Normal,
301 };
302 let response = handler.handle(&message).expect("operation should succeed");
303 assert!(response.success);
304 assert_eq!(response.message_id, "test_msg");
305 }
306 #[test]
307 fn test_consensus_manager() {
308 let mut consensus = ConsensusManager::new(
309 "node1".to_string(),
310 vec!["node2".to_string(), "node3".to_string()],
311 );
312 let (state, term) = consensus.get_state();
313 assert_eq!(state, ConsensusState::Follower);
314 assert_eq!(term, 0);
315 assert!(consensus.start_election().is_ok());
316 let (state, term) = consensus.get_state();
317 assert_eq!(state, ConsensusState::Candidate);
318 assert_eq!(term, 1);
319 let vote_request = VoteRequest {
320 term: 2,
321 candidate_id: "node2".to_string(),
322 last_log_index: 0,
323 last_log_term: 0,
324 };
325 let response = consensus.handle_vote_request(vote_request);
326 assert!(response.vote_granted);
327 assert_eq!(response.term, 2);
328 }
329 #[test]
330 fn test_consensus_log_entry() {
331 let mut consensus = ConsensusManager::new("node1".to_string(), vec!["node2".to_string()]);
332 let entry = LogEntry {
333 term: 1,
334 index: 0,
335 command: "test_command".to_string(),
336 data: vec![1, 2, 3],
337 };
338 assert!(consensus.append_entry(entry.clone()).is_err());
339 consensus.state = ConsensusState::Leader;
340 assert!(consensus.append_entry(entry).is_ok());
341 assert_eq!(consensus.log.len(), 1);
342 }
343 #[test]
344 fn test_data_partitioner() {
345 let mut partitioner = DataPartitioner::new(PartitioningStrategy::Hash, 4, 2);
346 let partition1 = partitioner.get_partition("key1");
347 let partition2 = partitioner.get_partition("key2");
348 assert!(partition1 < 4);
349 assert!(partition2 < 4);
350 assert_eq!(partition1, partitioner.get_partition("key1"));
351 partitioner.assign_partition(0, "node1".to_string());
352 partitioner.assign_partition(1, "node2".to_string());
353 let nodes = partitioner.get_partition_nodes(0);
354 assert!(nodes.contains(&"node1".to_string()));
355 }
356 #[test]
357 fn test_data_partitioner_rebalancing() {
358 let mut partitioner = DataPartitioner::new(PartitioningStrategy::Hash, 4, 1);
359 let nodes = vec!["node1".to_string(), "node2".to_string()];
360 let result = partitioner.rebalance_partitions(&nodes);
361 assert_eq!(result.assignments_changed, 4);
362 assert_eq!(result.partitions_moved.len(), 0);
363 }
364 #[test]
365 fn test_partitioning_strategies() {
366 let hash_partitioner = DataPartitioner::new(PartitioningStrategy::Hash, 4, 1);
367 let range_partitioner = DataPartitioner::new(PartitioningStrategy::Range, 4, 1);
368 let random_partitioner = DataPartitioner::new(PartitioningStrategy::Random, 4, 1);
369 let key = "test_key";
370 let hash_partition = hash_partitioner.get_partition(key);
371 let range_partition = range_partitioner.get_partition(key);
372 let random_partition = random_partitioner.get_partition(key);
373 assert!(hash_partition < 4);
374 assert!(range_partition < 4);
375 assert!(random_partition < 4);
376 }
377 #[test]
378 fn test_advanced_job_scheduler() {
379 let mut scheduler = AdvancedJobScheduler::new();
380 let mut nodes = HashMap::new();
381 let node1 = create_test_node("node1");
382 let node2 = create_test_node("node2");
383 nodes.insert("node1".to_string(), node1);
384 nodes.insert("node2".to_string(), node2);
385 let jobs = vec![create_test_job("job1"), create_test_job("job2")];
386 let decisions = scheduler
387 .gang_schedule(&jobs, &nodes)
388 .expect("operation should succeed");
389 assert_eq!(decisions.len(), jobs.len());
390 for decision in &decisions {
391 assert!(nodes.contains_key(&decision.node_id));
392 assert!(decision.resource_allocation.cpu_cores > 0);
393 }
394 }
395 #[test]
396 fn test_backfill_scheduling() {
397 let mut scheduler = AdvancedJobScheduler::new();
398 let mut nodes = HashMap::new();
399 let node1 = create_test_node("node1");
400 nodes.insert("node1".to_string(), node1);
401 let waiting_jobs = vec![create_test_job("waiting_job")];
402 let decisions = scheduler
403 .backfill_schedule(&waiting_jobs, &nodes)
404 .expect("operation should succeed");
405 assert_eq!(decisions.len(), 1);
406 assert_eq!(decisions[0].job_id, "waiting_job");
407 }
408 #[test]
409 fn test_resource_reservation() {
410 let mut scheduler = AdvancedJobScheduler::new();
411 let reservation = ResourceReservation {
412 id: "reservation1".to_string(),
413 node_id: "node1".to_string(),
414 start_time: Instant::now(),
415 duration: Duration::from_secs(3600),
416 resources: ResourceAllocation {
417 cpu_cores: 4,
418 memory_gb: 8,
419 gpu_count: 1,
420 storage_gb: 100,
421 network_bandwidth: 1000,
422 },
423 };
424 assert!(scheduler.reserve_resources(reservation).is_ok());
425 assert_eq!(scheduler.resource_reservations.len(), 1);
426 }
427 #[test]
428 fn test_checkpoint_manager() {
429 let mut checkpoint_mgr = CheckpointManager::new(Duration::from_secs(300));
430 let job_state = JobState {
431 progress: 0.5,
432 intermediate_results: HashMap::new(),
433 runtime_state: vec![1, 2, 3, 4],
434 };
435 let checkpoint_id = checkpoint_mgr
436 .create_checkpoint("job1", job_state.clone())
437 .expect("operation should succeed");
438 assert!(!checkpoint_id.is_empty());
439 let restored_state = checkpoint_mgr
440 .restore_checkpoint(&checkpoint_id)
441 .expect("operation should succeed");
442 assert_eq!(restored_state.progress, 0.5);
443 assert_eq!(restored_state.runtime_state, vec![1, 2, 3, 4]);
444 let stats = checkpoint_mgr.get_checkpoint_stats();
445 assert_eq!(stats.total_checkpoints, 1);
446 assert!(stats.total_size_bytes > 0);
447 }
448 #[test]
449 fn test_checkpoint_cleanup() {
450 let mut checkpoint_mgr = CheckpointManager::new(Duration::from_secs(300));
451 let job_state = JobState {
452 progress: 1.0,
453 intermediate_results: HashMap::new(),
454 runtime_state: vec![],
455 };
456 checkpoint_mgr
457 .create_checkpoint("job1", job_state.clone())
458 .expect("operation should succeed");
459 checkpoint_mgr
460 .create_checkpoint("job2", job_state)
461 .expect("operation should succeed");
462 assert_eq!(checkpoint_mgr.checkpoint_storage.len(), 2);
463 checkpoint_mgr.cleanup_old_checkpoints(Duration::from_secs(0));
464 assert_eq!(checkpoint_mgr.checkpoint_storage.len(), 0);
465 }
466 #[test]
467 fn test_message_type_conversion() {
468 assert_eq!(format!("{}", MessageType::JobSubmission), "job_submission");
469 assert_eq!(format!("{}", MessageType::JobResult), "job_result");
470 assert_eq!(format!("{}", MessageType::Heartbeat), "heartbeat");
471 assert_eq!(
472 format!("{}", MessageType::ResourceUpdate),
473 "resource_update"
474 );
475 assert_eq!(
476 format!("{}", MessageType::ConsensusRequest),
477 "consensus_request"
478 );
479 assert_eq!(format!("{}", MessageType::DataPartition), "data_partition");
480 assert_eq!(
481 format!("{}", MessageType::Custom("test".to_string())),
482 "test"
483 );
484 }
485 #[test]
486 fn test_consensus_states() {
487 let follower = ConsensusState::Follower;
488 let candidate = ConsensusState::Candidate;
489 let leader = ConsensusState::Leader;
490 assert_eq!(follower, ConsensusState::Follower);
491 assert_eq!(candidate, ConsensusState::Candidate);
492 assert_eq!(leader, ConsensusState::Leader);
493 }
494 #[test]
495 fn test_scheduling_policies() {
496 let policies = [
497 SchedulingPolicy::FIFO,
498 SchedulingPolicy::ShortestJobFirst,
499 SchedulingPolicy::GangScheduling,
500 SchedulingPolicy::Backfill,
501 SchedulingPolicy::PriorityBased,
502 ];
503 assert_eq!(policies.len(), 5);
504 }
505 #[test]
506 fn test_message_priorities() {
507 let low = MessagePriority::Low;
508 let normal = MessagePriority::Normal;
509 let high = MessagePriority::High;
510 let critical = MessagePriority::Critical;
511 match low {
512 MessagePriority::Low => {}
513 _ => panic!(),
514 }
515 match normal {
516 MessagePriority::Normal => {}
517 _ => panic!(),
518 }
519 match high {
520 MessagePriority::High => {}
521 _ => panic!(),
522 }
523 match critical {
524 MessagePriority::Critical => {}
525 _ => panic!(),
526 }
527 }
528 #[test]
529 fn test_resource_allocation_calculations() {
530 let allocation = ResourceAllocation {
531 cpu_cores: 8,
532 memory_gb: 16,
533 gpu_count: 2,
534 storage_gb: 500,
535 network_bandwidth: 1000,
536 };
537 assert_eq!(allocation.cpu_cores, 8);
538 assert_eq!(allocation.memory_gb, 16);
539 assert_eq!(allocation.gpu_count, 2);
540 assert_eq!(allocation.storage_gb, 500);
541 assert_eq!(allocation.network_bandwidth, 1000);
542 }
543 #[test]
544 fn test_job_state_serialization() {
545 let mut intermediate_results = HashMap::new();
546 intermediate_results.insert("result1".to_string(), vec![1, 2, 3]);
547 intermediate_results.insert("result2".to_string(), vec![4, 5, 6]);
548 let job_state = JobState {
549 progress: 0.75,
550 intermediate_results,
551 runtime_state: vec![7, 8, 9],
552 };
553 assert_eq!(job_state.progress, 0.75);
554 assert_eq!(job_state.intermediate_results.len(), 2);
555 assert_eq!(job_state.runtime_state, vec![7, 8, 9]);
556 assert!(job_state.intermediate_results.contains_key("result1"));
557 assert!(job_state.intermediate_results.contains_key("result2"));
558 }
559}