github.com/andy2046/gopie@v0.7.0/pkg/ratelimit/ratelimit_test.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  )
     7  
     8  var (
     9  	t0 = time.Date(2017, time.August, 8, 8, 8, 8, 0, time.UTC)
    10  )
    11  
    12  func TestEvery(t *testing.T) {
    13  	defer func() {
    14  		if r := recover(); r != nil {
    15  			if r != "ratelimit: invalid time interval for Every" {
    16  				t.Error("Every panic messgae error")
    17  			}
    18  		}
    19  	}()
    20  
    21  	cases := []struct {
    22  		interval time.Duration
    23  		l        Limit
    24  	}{
    25  		{1 * time.Second, Limit(1)},
    26  		{2 * time.Second, Limit(0.5)},
    27  		{4 * time.Second, Limit(0.25)},
    28  		{10 * time.Second, Limit(0.1)},
    29  		{1 * time.Millisecond, Limit(1e3)},
    30  		{10 * time.Millisecond, Limit(100)},
    31  		{-1 * time.Millisecond, Limit(-1)},
    32  	}
    33  	for _, c := range cases {
    34  		l := Every(c.interval)
    35  		if l-c.l > 0.0001 {
    36  			t.Errorf("Every(%v) = %v want %v", c.interval, l, c.l)
    37  		}
    38  	}
    39  }
    40  
    41  func TestAllow(t *testing.T) {
    42  	t1 := t0.Add(100 * time.Millisecond)
    43  	t2 := t0.Add(200 * time.Millisecond)
    44  	cases := []struct {
    45  		t  time.Time
    46  		n  int
    47  		ok bool
    48  	}{
    49  		{t0, 1, true},
    50  		{t0, 1, false},
    51  		{t0, 1, false},
    52  		{t1, 1, true},
    53  		{t1, 1, false},
    54  		{t1, 1, false},
    55  		{t2, 2, false}, // exceeds burst
    56  		{t2, 1, true},
    57  		{t2, 1, false},
    58  	}
    59  
    60  	l := New(10, 1)
    61  	for _, c := range cases {
    62  		ok := l.AllowN(c.t, c.n)
    63  		if ok != c.ok {
    64  			t.Errorf("AllowN(%v, %v) = %v want %v",
    65  				c.t, c.n, ok, c.ok)
    66  		}
    67  	}
    68  }
    69  
    70  func TestWait(t *testing.T) {
    71  	cases := []struct {
    72  		t      time.Time
    73  		n      int
    74  		nilErr bool
    75  	}{
    76  		{t0, 1, true},
    77  		{t0, 2, false}, // exceeds burst
    78  		{t0, 1, true},
    79  		{t0, 1, true},
    80  	}
    81  
    82  	l := New(10, 1)
    83  	for _, c := range cases {
    84  		w, err := l.WaitN(c.t, c.n)
    85  		delay := l.limit.durationFromTokens(float64(c.n))
    86  		if (c.nilErr && err != nil) || (!c.nilErr && err == nil) || w > delay {
    87  			t.Errorf("WaitN(%v, %v) = %v, %v want %v, nilErr? %v",
    88  				c.t, c.n, w, err, delay, c.nilErr)
    89  		}
    90  	}
    91  }