status-go/mailserver/limiter_test.go

110 lines
2.7 KiB
Go
Raw Permalink Normal View History

2018-05-21 11:30:37 +00:00
package mailserver
import (
"fmt"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestIsAllowed(t *testing.T) {
peerID := "peerID"
testCases := []struct {
t time.Duration
shouldBeAllowed bool
db func() map[string]time.Time
info string
}{
{
t: 5 * time.Millisecond,
shouldBeAllowed: true,
db: func() map[string]time.Time {
return make(map[string]time.Time)
},
info: "Expecting limiter.isAllowed to allow with an empty db",
2018-05-21 11:30:37 +00:00
},
{
t: 5 * time.Millisecond,
shouldBeAllowed: true,
db: func() map[string]time.Time {
db := make(map[string]time.Time)
db[peerID] = time.Now().Add(time.Duration(-10) * time.Millisecond)
return db
},
info: "Expecting limiter.isAllowed to allow with an expired peer on its db",
2018-05-21 11:30:37 +00:00
},
{
t: 5 * time.Millisecond,
shouldBeAllowed: false,
db: func() map[string]time.Time {
db := make(map[string]time.Time)
db[peerID] = time.Now().Add(time.Duration(-1) * time.Millisecond)
return db
},
info: "Expecting limiter.isAllowed to not allow with a non expired peer on its db",
2018-05-21 11:30:37 +00:00
},
}
for _, tc := range testCases {
t.Run(tc.info, func(*testing.T) {
l := newRateLimiter(tc.t)
2018-05-21 11:30:37 +00:00
l.db = tc.db()
assert.Equal(t, tc.shouldBeAllowed, l.IsAllowed(peerID))
2018-05-21 11:30:37 +00:00
})
}
}
func TestRemoveExpiredRateLimits(t *testing.T) {
peer := "peer"
l := newRateLimiter(time.Duration(5) * time.Second)
2018-05-21 11:30:37 +00:00
for i := 0; i < 10; i++ {
peerID := fmt.Sprintf("%s%d", peer, i)
l.db[peerID] = time.Now().Add(time.Duration(i*(-2)) * time.Second)
}
l.deleteExpired()
assert.Equal(t, 3, len(l.db))
for i := 0; i < 3; i++ {
peerID := fmt.Sprintf("%s%d", peer, i)
_, ok := l.db[peerID]
assert.True(t, ok, fmt.Sprintf("Non expired peer '%s' should exist, but it doesn't", peerID))
}
for i := 3; i < 10; i++ {
peerID := fmt.Sprintf("%s%d", peer, i)
_, ok := l.db[peerID]
assert.False(t, ok, fmt.Sprintf("Expired peer '%s' should not exist, but it does", peerID))
}
}
func TestCleaningUpExpiredRateLimits(t *testing.T) {
l := newRateLimiter(5 * time.Second)
l.period = time.Millisecond * 10
l.Start()
defer l.Stop()
l.db["peer01"] = time.Now().Add(-1 * time.Second)
l.db["peer02"] = time.Now().Add(-2 * time.Second)
l.db["peer03"] = time.Now().Add(-10 * time.Second)
time.Sleep(time.Millisecond * 20)
_, ok := l.db["peer01"]
assert.True(t, ok)
_, ok = l.db["peer02"]
assert.True(t, ok)
_, ok = l.db["peer03"]
assert.False(t, ok)
}
2018-05-21 11:30:37 +00:00
func TestAddingLimts(t *testing.T) {
peerID := "peerAdding"
l := newRateLimiter(time.Duration(5) * time.Second)
2018-05-21 11:30:37 +00:00
pre := time.Now()
l.Add(peerID)
2018-05-21 11:30:37 +00:00
post := time.Now()
assert.True(t, l.db[peerID].After(pre))
assert.True(t, l.db[peerID].Before(post))
}