diff --git a/crates/prune/types/src/limiter.rs b/crates/prune/types/src/limiter.rs index 09f0c4cd3..3a1059949 100644 --- a/crates/prune/types/src/limiter.rs +++ b/crates/prune/types/src/limiter.rs @@ -120,3 +120,270 @@ impl PruneLimiter { self.is_deleted_entries_limit_reached() || self.is_time_limit_reached() } } + +#[cfg(test)] +mod tests { + use super::*; + use std::thread::sleep; + + #[test] + fn test_prune_deleted_entries_limit_initial_state() { + let limit_tracker = PruneDeletedEntriesLimit::new(10); + // Limit should be set properly + assert_eq!(limit_tracker.limit, 10); + // No entries should be deleted + assert_eq!(limit_tracker.deleted, 0); + assert!(!limit_tracker.is_limit_reached()); + } + + #[test] + fn test_prune_deleted_entries_limit_is_limit_reached() { + // Test when the deleted entries are less than the limit + let mut limit_tracker = PruneDeletedEntriesLimit::new(5); + limit_tracker.deleted = 3; + assert!(!limit_tracker.is_limit_reached()); + + // Test when the deleted entries are equal to the limit + limit_tracker.deleted = 5; + assert!(limit_tracker.is_limit_reached()); + + // Test when the deleted entries exceed the limit + limit_tracker.deleted = 6; + assert!(limit_tracker.is_limit_reached()); + } + + #[test] + fn test_prune_time_limit_initial_state() { + let time_limit = PruneTimeLimit::new(Duration::from_secs(10)); + // The limit should be set correctly + assert_eq!(time_limit.limit, Duration::from_secs(10)); + // The elapsed time should be very small right after creation + assert!(time_limit.start.elapsed() < Duration::from_secs(1)); + // Limit should not be reached initially + assert!(!time_limit.is_limit_reached()); + } + + #[test] + fn test_prune_time_limit_is_limit_reached() { + let time_limit = PruneTimeLimit::new(Duration::from_millis(50)); + + // Simulate waiting for some time (less than the limit) + std::thread::sleep(Duration::from_millis(30)); + assert!(!time_limit.is_limit_reached()); + + // Simulate waiting for time greater than the limit + std::thread::sleep(Duration::from_millis(30)); + assert!(time_limit.is_limit_reached()); + } + + #[test] + fn test_set_deleted_entries_limit_initial_state() { + let pruner = PruneLimiter::default().set_deleted_entries_limit(100); + // The deleted_entries_limit should be set with the correct limit + assert!(pruner.deleted_entries_limit.is_some()); + let deleted_entries_limit = pruner.deleted_entries_limit.unwrap(); + assert_eq!(deleted_entries_limit.limit, 100); + // The deleted count should be initially zero + assert_eq!(deleted_entries_limit.deleted, 0); + // The limit should not be reached initially + assert!(!deleted_entries_limit.is_limit_reached()); + } + + #[test] + fn test_set_deleted_entries_limit_overwrite_existing() { + let mut pruner = PruneLimiter::default().set_deleted_entries_limit(50); + // Overwrite the existing limit + pruner = pruner.set_deleted_entries_limit(200); + + assert!(pruner.deleted_entries_limit.is_some()); + let deleted_entries_limit = pruner.deleted_entries_limit.unwrap(); + // Check that the limit has been overwritten correctly + assert_eq!(deleted_entries_limit.limit, 200); + // Deleted count should still be zero + assert_eq!(deleted_entries_limit.deleted, 0); + assert!(!deleted_entries_limit.is_limit_reached()); + } + + #[test] + fn test_set_deleted_entries_limit_when_limit_is_reached() { + let mut pruner = PruneLimiter::default().set_deleted_entries_limit(5); + assert!(pruner.deleted_entries_limit.is_some()); + let mut deleted_entries_limit = pruner.deleted_entries_limit.clone().unwrap(); + + // Simulate deletion of entries + deleted_entries_limit.deleted = 5; + assert!(deleted_entries_limit.is_limit_reached()); + + // Overwrite the limit and check if it resets correctly + pruner = pruner.set_deleted_entries_limit(10); + deleted_entries_limit = pruner.deleted_entries_limit.unwrap(); + assert_eq!(deleted_entries_limit.limit, 10); + // Deletion count should reset + assert_eq!(deleted_entries_limit.deleted, 0); + assert!(!deleted_entries_limit.is_limit_reached()); + } + + #[test] + fn test_floor_deleted_entries_limit_to_multiple_of() { + let limiter = PruneLimiter::default().set_deleted_entries_limit(15); + let denominator = NonZeroUsize::new(4).unwrap(); + + // Floor limit to the largest multiple of 4 less than or equal to 15 (that is 12) + let updated_limiter = limiter.floor_deleted_entries_limit_to_multiple_of(denominator); + assert_eq!(updated_limiter.deleted_entries_limit.unwrap().limit, 12); + + // Test when the limit is already a multiple of the denominator + let limiter = PruneLimiter::default().set_deleted_entries_limit(16); + let updated_limiter = limiter.floor_deleted_entries_limit_to_multiple_of(denominator); + assert_eq!(updated_limiter.deleted_entries_limit.unwrap().limit, 16); + + // Test when there's no limit set (should not panic) + let limiter = PruneLimiter::default(); + let updated_limiter = limiter.floor_deleted_entries_limit_to_multiple_of(denominator); + assert!(updated_limiter.deleted_entries_limit.is_none()); + } + + #[test] + fn test_is_deleted_entries_limit_reached() { + // Limit is not set, should return false + let limiter = PruneLimiter::default(); + assert!(!limiter.is_deleted_entries_limit_reached()); + + // Limit is set but not reached, should return false + let mut limiter = PruneLimiter::default().set_deleted_entries_limit(10); + limiter.deleted_entries_limit.as_mut().unwrap().deleted = 5; + // 5 entries deleted out of 10 + assert!(!limiter.is_deleted_entries_limit_reached()); + + // Limit is reached, should return true + limiter.deleted_entries_limit.as_mut().unwrap().deleted = 10; + // 10 entries deleted out of 10 + assert!(limiter.is_deleted_entries_limit_reached()); + + // Deleted entries exceed the limit, should return true + limiter.deleted_entries_limit.as_mut().unwrap().deleted = 12; + // 12 entries deleted out of 10 + assert!(limiter.is_deleted_entries_limit_reached()); + } + + #[test] + fn test_increment_deleted_entries_count_by() { + // Increment when no limit is set + let mut limiter = PruneLimiter::default(); + limiter.increment_deleted_entries_count_by(5); + assert_eq!(limiter.deleted_entries_limit.as_ref().map(|l| l.deleted), None); // Still None + + // Increment when limit is set + let mut limiter = PruneLimiter::default().set_deleted_entries_limit(10); + limiter.increment_deleted_entries_count_by(3); + assert_eq!(limiter.deleted_entries_limit.as_ref().unwrap().deleted, 3); // Now 3 deleted + + // Increment again + limiter.increment_deleted_entries_count_by(2); + assert_eq!(limiter.deleted_entries_limit.as_ref().unwrap().deleted, 5); // Now 5 deleted + } + + #[test] + fn test_increment_deleted_entries_count() { + let mut limiter = PruneLimiter::default().set_deleted_entries_limit(5); + assert_eq!(limiter.deleted_entries_limit.as_ref().unwrap().deleted, 0); // Initially 0 + + limiter.increment_deleted_entries_count(); // Increment by 1 + assert_eq!(limiter.deleted_entries_limit.as_ref().unwrap().deleted, 1); // Now 1 + } + + #[test] + fn test_deleted_entries_limit_left() { + // Test when limit is set and some entries are deleted + let mut limiter = PruneLimiter::default().set_deleted_entries_limit(10); + limiter.increment_deleted_entries_count_by(3); // Simulate 3 deleted entries + assert_eq!(limiter.deleted_entries_limit_left(), Some(7)); // 10 - 3 = 7 + + // Test when no entries are deleted + limiter = PruneLimiter::default().set_deleted_entries_limit(5); + assert_eq!(limiter.deleted_entries_limit_left(), Some(5)); // 5 - 0 = 5 + + // Test when limit is reached + limiter.increment_deleted_entries_count_by(5); // Simulate deleting 5 entries + assert_eq!(limiter.deleted_entries_limit_left(), Some(0)); // 5 - 5 = 0 + + // Test when limit is not set + limiter = PruneLimiter::default(); // No limit set + assert_eq!(limiter.deleted_entries_limit_left(), None); // Should be None + } + + #[test] + fn test_set_time_limit() { + // Create a PruneLimiter instance with no time limit set + let mut limiter = PruneLimiter::default(); + + // Set a time limit of 5 seconds + limiter = limiter.set_time_limit(Duration::new(5, 0)); + + // Verify that the time limit is set correctly + assert!(limiter.time_limit.is_some()); + let time_limit = limiter.time_limit.as_ref().unwrap(); + assert_eq!(time_limit.limit, Duration::new(5, 0)); + // Ensure the start time is recent + assert!(time_limit.start.elapsed() < Duration::new(1, 0)); + } + + #[test] + fn test_is_time_limit_reached() { + // Create a PruneLimiter instance and set a time limit of 10 milliseconds + let mut limiter = PruneLimiter::default(); + + // Time limit should not be reached initially + assert!(!limiter.is_time_limit_reached(), "Time limit should not be reached yet"); + + limiter = limiter.set_time_limit(Duration::new(0, 10_000_000)); // 10 milliseconds + + // Sleep for 5 milliseconds (less than the time limit) + sleep(Duration::new(0, 5_000_000)); // 5 milliseconds + assert!(!limiter.is_time_limit_reached(), "Time limit should not be reached yet"); + + // Sleep for an additional 10 milliseconds (totaling 15 milliseconds) + sleep(Duration::new(0, 10_000_000)); // 10 milliseconds + assert!(limiter.is_time_limit_reached(), "Time limit should be reached now"); + } + + #[test] + fn test_is_limit_reached() { + // Create a PruneLimiter instance + let mut limiter = PruneLimiter::default(); + + // Test when no limits are set + assert!(!limiter.is_limit_reached(), "Limit should not be reached with no limits set"); + + // Set a deleted entries limit + limiter = limiter.set_deleted_entries_limit(5); + assert!( + !limiter.is_limit_reached(), + "Limit should not be reached when deleted entries are less than limit" + ); + + // Increment deleted entries count to reach the limit + limiter.increment_deleted_entries_count_by(5); + assert!( + limiter.is_limit_reached(), + "Limit should be reached when deleted entries equal the limit" + ); + + // Reset the limiter + limiter = PruneLimiter::default(); + + // Set a time limit and check + limiter = limiter.set_time_limit(Duration::new(0, 10_000_000)); // 10 milliseconds + + // Sleep for 5 milliseconds (less than the time limit) + sleep(Duration::new(0, 5_000_000)); // 5 milliseconds + assert!( + !limiter.is_limit_reached(), + "Limit should not be reached when time limit not reached" + ); + + // Sleep for another 10 milliseconds (totaling 15 milliseconds) + sleep(Duration::new(0, 10_000_000)); // 10 milliseconds + assert!(limiter.is_limit_reached(), "Limit should be reached when time limit is reached"); + } +}