woop/internal/rate/limiter_test.go

package rate

import (
	"container/list"
	"testing"
	"time"

	"golang.org/x/time/rate"
)

var (
	id1 = "id1"
	id2 = "id2"
)

func TestLimiterPerID_AllowN(t *testing.T) {
	tests := []struct {
		steps []testStep
	}{
		{
			steps: []testStep{
				{id1, 1, true},
				{id1, 1, true},
				{id1, 1, false},
			},
		},
		{
			steps: []testStep{
				{id1, 1, true},
				{id2, 1, true},
				{id1, 1, true},
				{id1, 1, false},
				{id2, 1, true},
			},
		},
		{
			steps: []testStep{
				{id1, 2, true},
				{id1, 1, false},
			},
		},
		{
			steps: []testStep{
				{id1, 3, false},
				{id1, 3, false},
			},
		},
	}
	for i, test := range tests {
		lpi := newTestLimiter()
		steps := test.steps
		for j, step := range steps {
			res := lpi.AllowN(step.id, step.n)
			if res != step.exp {
				t.Errorf("Test %v/%v unexpected %v/%v", i, j, res, step.exp)
			}
			lpi.t.(*testTimer).tick()
		}
	}
}

func TestLimiterPerID_maintain(t *testing.T) {
	lpi := newTestLimiter()
	lpi.c.capacity = 0
	lpi.AllowN(id1, 2)
	lpi.t.(*testTimer).tick()
	lpi.AllowN(id2, 2)
	lpi.t.(*testTimer).tick()

	lpi.maintain()
	if lpi.evictList.Len() != 1 || len(lpi.items) != 1 {
		t.Errorf("unexpected number. Expect 1")
	}
	for id := range lpi.items {
		if id != id2 {
			t.Errorf("unexpected id. Expect id2")
		}
	}
	lpi.t.(*testTimer).tick()
	lpi.maintain()
	if lpi.evictList.Len() != 0 || len(lpi.items) != 0 {
		t.Errorf("unexpected number. Expect 0")
	}
}

func TestLimiterPerID_maintain_largeCap(t *testing.T) {
	lpi := newTestLimiter()
	lpi.c.capacity = 10
	lpi.AllowN(id1, 2)
	lpi.t.(*testTimer).tick()
	lpi.AllowN(id2, 2)
	lpi.t.(*testTimer).tick()
	lpi.maintain()

	if lpi.evictList.Len() != 2 || len(lpi.items) != 2 {
		t.Errorf("unexpected number")
	}
}

func TestLimiterPerID_maintain_largeMinDur(t *testing.T) {
	lpi := newTestLimiter()
	lpi.c.minEvictDur = 5 * time.Second
	lpi.AllowN(id1, 2)
	lpi.t.(*testTimer).tick()
	lpi.AllowN(id2, 2)
	lpi.t.(*testTimer).tick()
	lpi.maintain()

	if lpi.evictList.Len() != 2 || len(lpi.items) != 2 {
		t.Errorf("unexpected number")
	}
}

type testStep struct {
	id  string
	n   int
	exp bool
}

func newTestLimiter() *limiterPerID {
	return &limiterPerID{
		evictList: list.New(),
		items:     make(map[string]*list.Element),
		t:         &testTimer{time.Now()},
		c:         testConfig,
	}
}

var testConfig = configInt{
	limit:       rate.Every(100 * time.Second),
	burst:       2,
	capacity:    1,
	checkInt:    time.Second,
	minEvictDur: 1 * time.Second,
	whitelist:   make(map[string]struct{}),
}

// testTimer will increment 1 sec per call
type testTimer struct {
	c time.Time
}

func (t *testTimer) now() time.Time {
	return t.c
}

func (t *testTimer) tick() {
	t.c = t.c.Add(time.Second)
}

func (t *testTimer) since(t2 time.Time) time.Duration {
	return t.c.Sub(t2)
}

func (t *testTimer) newTicker(d time.Duration) *time.Ticker {
	return time.NewTicker(time.Nanosecond)
}