Add unit tests for tokenbucket.py

Author: dmptrluke

cloudbot/util/test/test_tokenbucket.py

@@ -0,0 +1,42 @@
+import time
+
+from cloudbot.util.tokenbucket import TokenBucket
+
+
+# noinspection PyProtectedMember
+def test_bucket_consume():
+    bucket = TokenBucket(10, 5)
+    # larger then capacity
+    assert bucket.consume(15) is False
+    # success
+    assert bucket.consume(10) is True
+    # check if bucket has no tokens
+    assert bucket._tokens == 0
+    # bucket is empty from above, should fail
+    assert bucket.consume(10) is False
+
+
+# noinspection PyProtectedMember
+def test_bucket_advanced():
+    bucket = TokenBucket(10, 1)
+    # tokens start at 10
+    assert bucket._tokens == 10
+    # empty tokens
+    assert bucket.empty() is True
+    # check tokens is 0
+    assert bucket._tokens == 0
+    # refill tokens
+    assert bucket.refill() is True
+    # check tokens is 10
+    assert bucket._tokens == 10
+
+
+def test_bucket_regen():
+    bucket = TokenBucket(10, 10)
+    # success
+    assert bucket.consume(10) is True
+    # sleep
+    time.sleep(1)
+    # bucket should be full again and this should succeed
+    assert bucket.tokens == 10
+    assert bucket.consume(10) is True

cloudbot/util/tokenbucket.py

@@ -11,30 +11,51 @@ class TokenBucket(object):
     False
     """
 
-    def __init__(self, tokens, fill_rate):
+    def __init__(self, _capacity, fill_rate):
+        """
+        :param _capacity: The total amount of token the bucket can contain
+        :param fill_rate: The rate at which tokens regenerate. (fill_rate per second)
+        """
         """tokens is the total tokens in the bucket. fill_rate is the
         rate in tokens/second that the bucket will be refilled."""
-        self.capacity = float(tokens)
-        self._tokens = float(tokens)
+        self.capacity = float(_capacity)
+        self._tokens = float(_capacity)
         self.fill_rate = float(fill_rate)
         self.timestamp = time()
 
     def consume(self, tokens):
-        """Consume tokens from the bucket. Returns True if there were
-        sufficient tokens otherwise False."""
+        """
+        Consume tokens from the bucket.
+        :param tokens: The number of tokens to consume
+        :return true if there were sufficient tokens otherwise false
+        """
         if tokens <= self.tokens:
             self._tokens -= tokens
         else:
             return False
         return True
 
     def refill(self):
+        """
+        Sets the current token count to the max capacity
+        """
         self._tokens = self.capacity
+        return True
 
     def empty(self):
+        """
+        Sets the current token count to zero
+        """
         self._tokens = float(0)
+        return True
 
     def get_tokens(self):
+        """
+        Calculates and returns the current amount of tokens the bucker contains
+
+        :return Amount of tokens the bucket contains
+        :rtype Float
+        """
         now = time()
         if self._tokens < self.capacity:
             delta = self.fill_rate * (now - self.timestamp)