gopkg.in/hpcloud/tail.v1@v1.0.0/ratelimiter/leakybucket_test.go (about)

     1  package ratelimiter
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  )
     7  
     8  func TestPour(t *testing.T) {
     9  	bucket := NewLeakyBucket(60, time.Second)
    10  	bucket.Lastupdate = time.Unix(0, 0)
    11  
    12  	bucket.Now = func() time.Time { return time.Unix(1, 0) }
    13  
    14  	if bucket.Pour(61) {
    15  		t.Error("Expected false")
    16  	}
    17  
    18  	if !bucket.Pour(10) {
    19  		t.Error("Expected true")
    20  	}
    21  
    22  	if !bucket.Pour(49) {
    23  		t.Error("Expected true")
    24  	}
    25  
    26  	if bucket.Pour(2) {
    27  		t.Error("Expected false")
    28  	}
    29  
    30  	bucket.Now = func() time.Time { return time.Unix(61, 0) }
    31  	if !bucket.Pour(60) {
    32  		t.Error("Expected true")
    33  	}
    34  
    35  	if bucket.Pour(1) {
    36  		t.Error("Expected false")
    37  	}
    38  
    39  	bucket.Now = func() time.Time { return time.Unix(70, 0) }
    40  
    41  	if !bucket.Pour(1) {
    42  		t.Error("Expected true")
    43  	}
    44  
    45  }
    46  
    47  func TestTimeSinceLastUpdate(t *testing.T) {
    48  	bucket := NewLeakyBucket(60, time.Second)
    49  	bucket.Now = func() time.Time { return time.Unix(1, 0) }
    50  	bucket.Pour(1)
    51  	bucket.Now = func() time.Time { return time.Unix(2, 0) }
    52  
    53  	sinceLast := bucket.TimeSinceLastUpdate()
    54  	if sinceLast != time.Second*1 {
    55  		t.Errorf("Expected time since last update to be less than 1 second, got %d", sinceLast)
    56  	}
    57  }
    58  
    59  func TestTimeToDrain(t *testing.T) {
    60  	bucket := NewLeakyBucket(60, time.Second)
    61  	bucket.Now = func() time.Time { return time.Unix(1, 0) }
    62  	bucket.Pour(10)
    63  
    64  	if bucket.TimeToDrain() != time.Second*10 {
    65  		t.Error("Time to drain should be 10 seconds")
    66  	}
    67  
    68  	bucket.Now = func() time.Time { return time.Unix(2, 0) }
    69  
    70  	if bucket.TimeToDrain() != time.Second*9 {
    71  		t.Error("Time to drain should be 9 seconds")
    72  	}
    73  }