mirror of
https://github.com/empayre/fleet.git
synced 2024-11-06 17:05:18 +00:00
128 lines
3.3 KiB
Go
128 lines
3.3 KiB
Go
package redis_test
|
|
|
|
import (
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/fleetdm/fleet/v4/server/datastore/redis"
|
|
"github.com/fleetdm/fleet/v4/server/datastore/redis/redistest"
|
|
"github.com/fleetdm/fleet/v4/server/fleet"
|
|
redigo "github.com/gomodule/redigo/redis"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestThrottledStore(t *testing.T) {
|
|
const prefix = "TestThrottledStore:"
|
|
|
|
runTest := func(t *testing.T, pool fleet.RedisPool) {
|
|
store := redis.ThrottledStore{
|
|
Pool: pool,
|
|
KeyPrefix: prefix,
|
|
}
|
|
|
|
t.Run("GetWithTime", func(t *testing.T) {
|
|
conn := pool.Get()
|
|
defer conn.Close()
|
|
|
|
// key does not exist
|
|
v, ts, err := store.GetWithTime("a")
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, int64(-1))
|
|
require.WithinDuration(t, time.Now(), ts, time.Second)
|
|
|
|
_, err = conn.Do("SET", store.KeyPrefix+"a", 1)
|
|
require.NoError(t, err)
|
|
|
|
// key exists
|
|
v, ts, err = store.GetWithTime("a")
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, int64(1))
|
|
require.WithinDuration(t, time.Now(), ts, time.Second)
|
|
})
|
|
|
|
t.Run("SetIfNotExistsWithTTL", func(t *testing.T) {
|
|
conn := pool.Get()
|
|
defer conn.Close()
|
|
|
|
// key does not exist
|
|
ok, err := store.SetIfNotExistsWithTTL("{b}", 1, time.Second)
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
|
|
v, err := redigo.Int(conn.Do("GET", store.KeyPrefix+"{b}"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, 1)
|
|
|
|
// key exists
|
|
ok, err = store.SetIfNotExistsWithTTL("{b}", 2, time.Second)
|
|
require.NoError(t, err)
|
|
require.False(t, ok)
|
|
|
|
// value is still 1
|
|
v, err = redigo.Int(conn.Do("GET", store.KeyPrefix+"{b}"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, 1)
|
|
|
|
// key does not exist, but ttl less than a second
|
|
ok, err = store.SetIfNotExistsWithTTL("{b}2", 3, time.Millisecond)
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
|
|
v, err = redigo.Int(conn.Do("GET", store.KeyPrefix+"{b}2"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, 3)
|
|
})
|
|
|
|
t.Run("CompareAndSwapWithTTL", func(t *testing.T) {
|
|
conn := pool.Get()
|
|
defer conn.Close()
|
|
|
|
// key does not exist
|
|
ok, err := store.CompareAndSwapWithTTL("{c}", 1, 2, time.Second)
|
|
require.NoError(t, err)
|
|
require.False(t, ok)
|
|
|
|
_, err = conn.Do("SET", store.KeyPrefix+"{c}", 1)
|
|
require.NoError(t, err)
|
|
|
|
// key exists, but values do not match
|
|
ok, err = store.CompareAndSwapWithTTL("{c}", 2, 3, time.Second)
|
|
require.NoError(t, err)
|
|
require.False(t, ok)
|
|
|
|
// key exists, values match
|
|
ok, err = store.CompareAndSwapWithTTL("{c}", 1, 4, time.Second)
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
|
|
v, err := redigo.Int(conn.Do("GET", store.KeyPrefix+"{c}"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, 4)
|
|
|
|
// key exists, ttl less than a second
|
|
ok, err = store.CompareAndSwapWithTTL("{c}", 4, 5, time.Millisecond)
|
|
require.NoError(t, err)
|
|
require.True(t, ok)
|
|
|
|
v, err = redigo.Int(conn.Do("GET", store.KeyPrefix+"{c}"))
|
|
require.NoError(t, err)
|
|
require.Equal(t, v, 5)
|
|
})
|
|
}
|
|
|
|
t.Run("standalone", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, prefix, false, false, false)
|
|
runTest(t, pool)
|
|
})
|
|
|
|
t.Run("cluster", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, prefix, true, true, false)
|
|
runTest(t, pool)
|
|
})
|
|
|
|
t.Run("cluster_nofollow", func(t *testing.T) {
|
|
pool := redistest.SetupRedis(t, prefix, true, false, false)
|
|
runTest(t, pool)
|
|
})
|
|
}
|