110 lines
2.7 KiB
Go
110 lines
2.7 KiB
Go
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",
|
|
},
|
|
{
|
|
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",
|
|
},
|
|
{
|
|
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",
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.info, func(*testing.T) {
|
|
l := newRateLimiter(tc.t)
|
|
l.db = tc.db()
|
|
assert.Equal(t, tc.shouldBeAllowed, l.IsAllowed(peerID))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestRemoveExpiredRateLimits(t *testing.T) {
|
|
peer := "peer"
|
|
l := newRateLimiter(time.Duration(5) * time.Second)
|
|
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)
|
|
}
|
|
|
|
func TestAddingLimts(t *testing.T) {
|
|
peerID := "peerAdding"
|
|
l := newRateLimiter(time.Duration(5) * time.Second)
|
|
pre := time.Now()
|
|
l.Add(peerID)
|
|
post := time.Now()
|
|
assert.True(t, l.db[peerID].After(pre))
|
|
assert.True(t, l.db[peerID].Before(post))
|
|
}
|