mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
51 lines
1.5 KiB
Go
51 lines
1.5 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/alicebob/miniredis/v2"
|
|
"github.com/redis/go-redis/v9"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestOpenAIAPIKeyHealthCacheTripsWithinRollingWindow(t *testing.T) {
|
|
server := miniredis.RunT(t)
|
|
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
|
t.Cleanup(func() { _ = client.Close() })
|
|
store, ok := NewTempUnschedCache(client).(service.OpenAIAPIKeyHealthCache)
|
|
require.True(t, ok)
|
|
|
|
ctx := context.Background()
|
|
for attempt := 1; attempt <= 3; attempt++ {
|
|
count, tripped, err := store.RecordOpenAIAPIKeyHealthFailure(ctx, 42, 1, 3)
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, attempt, count)
|
|
require.Equal(t, attempt == 3, tripped)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIAPIKeyHealthCacheDropsFailuresOutsideRollingWindow(t *testing.T) {
|
|
server := miniredis.RunT(t)
|
|
now := time.Unix(1_700_000_000, 0)
|
|
server.SetTime(now)
|
|
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
|
t.Cleanup(func() { _ = client.Close() })
|
|
store, ok := NewTempUnschedCache(client).(service.OpenAIAPIKeyHealthCache)
|
|
require.True(t, ok)
|
|
|
|
ctx := context.Background()
|
|
count, tripped, err := store.RecordOpenAIAPIKeyHealthFailure(ctx, 42, 1, 3)
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, 1, count)
|
|
require.False(t, tripped)
|
|
|
|
server.SetTime(now.Add(61 * time.Second))
|
|
count, tripped, err = store.RecordOpenAIAPIKeyHealthFailure(ctx, 42, 1, 3)
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, 1, count)
|
|
require.False(t, tripped)
|
|
}
|