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