mirror of
https://github.com/binwiederhier/ntfy.git
synced 2026-11-10 14:44:04 +08:00
binwiederhier
2c91effa30
Nothing outside the limit package read it: enforcement is the bandwidth token bucket, and stats only report messages, emails and calls. Without the tally, a key is forgotten once all of its buckets are full again and it was idle for settledIdle, like a key that used no bandwidth. Usage now holds only the daily counts.
77 lines
1.9 KiB
Go
77 lines
1.9 KiB
Go
package limit
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestKind_String_ParseKind(t *testing.T) {
|
|
for k := Kind(0); k < numKinds; k++ {
|
|
parsed, ok := ParseKind(k.String())
|
|
require.True(t, ok, k.String())
|
|
require.Equal(t, k, parsed)
|
|
}
|
|
require.Equal(t, "messages", KindMessages.String())
|
|
require.Equal(t, "bandwidth", KindBandwidth.String())
|
|
_, ok := ParseKind("nope")
|
|
require.False(t, ok)
|
|
}
|
|
|
|
func TestUsage_Count(t *testing.T) {
|
|
s := newTestService(t)
|
|
require.True(t, s.Allow("k", Rule{
|
|
Kind: KindMessages,
|
|
Daily: 10,
|
|
}, 3).OK)
|
|
require.True(t, s.Allow("k", Rule{
|
|
Kind: KindBandwidth,
|
|
Burst: 100,
|
|
Refill: time.Second,
|
|
}, 42).OK)
|
|
require.True(t, s.Allow("k", Rule{
|
|
Kind: KindRequests,
|
|
Burst: 100,
|
|
Refill: time.Second,
|
|
}, 1).OK)
|
|
u := s.Usage("k")
|
|
require.Equal(t, int64(3), u.Count(KindMessages))
|
|
require.Equal(t, int64(0), u.Count(KindEmails))
|
|
require.Equal(t, int64(0), u.Count(KindBandwidth)) // No daily count, so that it does not hold the key (see Prune)
|
|
require.Equal(t, int64(0), u.Count(KindRequests))
|
|
}
|
|
|
|
func TestLimiter_Value_Bandwidth(t *testing.T) {
|
|
s := newTestService(t)
|
|
l := s.Limiter("k", Rule{
|
|
Kind: KindBandwidth,
|
|
Burst: 1000,
|
|
Refill: time.Second,
|
|
})
|
|
require.True(t, l.AllowN(300))
|
|
require.Zero(t, l.Value()) // No daily count; the bucket is what limits it
|
|
require.Equal(t, int64(700), s.Tokens("k", Rule{
|
|
Kind: KindBandwidth,
|
|
Burst: 1000,
|
|
Refill: time.Second,
|
|
}))
|
|
}
|
|
|
|
func TestLimiter_AllowNRefundsNegative(t *testing.T) {
|
|
s := newTestService(t)
|
|
l := s.Limiter("k", Rule{
|
|
Kind: KindMessages,
|
|
Burst: 10,
|
|
Refill: time.Hour,
|
|
})
|
|
require.True(t, l.AllowN(8))
|
|
require.False(t, l.AllowN(3))
|
|
require.Equal(t, int64(8), l.Value())
|
|
require.True(t, l.AllowN(-5)) // Refund, as the limit reader does when a later limiter rejects
|
|
require.True(t, l.AllowN(3))
|
|
require.Equal(t, int64(6), l.Value())
|
|
require.True(t, l.Allow())
|
|
require.Equal(t, int64(7), l.Value())
|
|
}
|