mirror of
https://github.com/siderolabs/discovery-service.git
synced 2025-03-14 09:55:08 +00:00
c55142668f
Update removing multiple old middlewares, rework the way data is passed through the context, logging fields, etc. Fix minimum keepalive interval enforcement. Signed-off-by: Andrey Smirnov <andrey.smirnov@siderolabs.com>
64 lines
1.3 KiB
Go
64 lines
1.3 KiB
Go
// Copyright (c) 2024 Sidero Labs, Inc.
|
|
//
|
|
// Use of this software is governed by the Business Source License
|
|
// included in the LICENSE file.
|
|
|
|
package limiter_test
|
|
|
|
import (
|
|
"net/netip"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/siderolabs/discovery-service/internal/limiter"
|
|
)
|
|
|
|
func TestDoGC(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
testIP1 = netip.MustParseAddr("1.1.1.1")
|
|
testIP2 = netip.MustParseAddr("2.2.2.2")
|
|
)
|
|
|
|
limiter := limiter.NewIPRateLimiter(1, 1)
|
|
|
|
// hit the bucket size
|
|
lim1 := limiter.Get(testIP1)
|
|
lim2 := limiter.Get(testIP2)
|
|
lim2.Allow()
|
|
|
|
assert.Exactly(t, 2, limiter.Len())
|
|
|
|
limiter.DoGC(time.Now())
|
|
|
|
assert.Exactly(t, 1, limiter.Len())
|
|
|
|
// lim1 should have been gc'd, but not lim2
|
|
assert.NotSame(t, lim1, limiter.Get(testIP1))
|
|
assert.Same(t, lim2, limiter.Get(testIP2))
|
|
|
|
// one more gc with T+2
|
|
limiter.DoGC(time.Now().Add(2 * time.Second))
|
|
|
|
assert.Exactly(t, 0, limiter.Len())
|
|
}
|
|
|
|
func TestIndependentLimiters(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
var (
|
|
testIP1 = netip.MustParseAddr("1.1.1.1")
|
|
testIP2 = netip.MustParseAddr("2.2.2.2")
|
|
)
|
|
|
|
limiter := limiter.NewIPRateLimiter(1, 1)
|
|
|
|
require.True(t, limiter.Get(testIP1).Allow())
|
|
require.False(t, limiter.Get(testIP1).Allow())
|
|
require.True(t, limiter.Get(testIP2).Allow())
|
|
}
|