From 4eb5c0fa61cddd297f49c3dfc35be9a27031d24d Mon Sep 17 00:00:00 2001
From: Jake Meyer <jakewmeyer@gmail.com>
Date: Fri, 24 Nov 2023 14:11:27 -0800
Subject: [PATCH] Add underflow/overflow tests to token bucket

---
 src/token_bucket.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/src/token_bucket.rs b/src/token_bucket.rs
index 770e3e6..b1cc5ad 100644
--- a/src/token_bucket.rs
+++ b/src/token_bucket.rs
@@ -52,6 +52,8 @@ impl TokenBucket {
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use super::*;
 
     #[test]
@@ -69,4 +71,31 @@ mod tests {
         assert!(bucket.take());
         assert!(!bucket.take());
     }
+
+    #[test]
+    fn test_token_bucket_fill_rate_overflow() {
+        let mut bucket = TokenBucket::new(255, 255, 1);
+        std::thread::sleep(std::time::Duration::from_secs(1));
+        assert!(bucket.take());
+        assert_eq!(bucket.available_tokens, 254);
+    }
+
+    #[test]
+    fn test_token_bucket_overflow() {
+        const PAST: u64 = 1_000_000_000;
+        let mut bucket = TokenBucket::new(255, 255, 1);
+        // Check that we can't overflow tokens_to_add
+        bucket.last_update = Instant::now() - Duration::from_secs(PAST);
+        assert!(bucket.take());
+    }
+
+    #[test]
+    fn test_token_bucket_underflow() {
+        const FUTURE: u64 = 1_000_000_000;
+        let mut bucket = TokenBucket::new(255, 1, 1);
+        // Check that we can't underflow tokens_to_add
+        bucket.last_update = Instant::now() + Duration::from_secs(FUTURE);
+        assert!(bucket.take());
+        assert_eq!(bucket.available_tokens, 254);
+    }
 }