111 lines
4.5 KiB
Go
111 lines
4.5 KiB
Go
package perfmetrics
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/alicebob/miniredis/v2"
|
|
"github.com/go-redis/redis/v8"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestLiveUserResults(t *testing.T) {
|
|
oldRedis := common.RedisEnabled
|
|
common.RedisEnabled = false
|
|
t.Cleanup(func() { common.RedisEnabled = oldRedis })
|
|
canceled, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
for _, tc := range []struct {
|
|
name string
|
|
test bool
|
|
token int
|
|
ctx context.Context
|
|
err *types.NewAPIError
|
|
want bool
|
|
success bool
|
|
}{
|
|
{"success", false, 1, context.Background(), nil, true, true},
|
|
{"upstream 503", false, 1, context.Background(), types.InitOpenAIError("server_error", 503), true, false},
|
|
{"probe", true, 1, context.Background(), types.InitOpenAIError("server_error", 503), false, false},
|
|
{"internal", false, 0, context.Background(), nil, false, false},
|
|
{"canceled", false, 1, canceled, types.InitOpenAIError("server_error", 503), false, false},
|
|
{"bad prompt", false, 1, context.Background(), types.InitOpenAIError("context_length_exceeded", 400), false, false},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Cleanup(func() { liveResults.Delete(tc.name) })
|
|
RecordRelayResult(tc.ctx, &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, IsStream: true, StartTime: time.Now()}, tc.err)
|
|
result := LatestUserResults(context.Background(), []string{tc.name})[tc.name]
|
|
assert.Equal(t, tc.want, result.At > 0)
|
|
if tc.want {
|
|
assert.Equal(t, tc.success, result.Success)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestUserResultsAreOrderedAndSurviveRestart(t *testing.T) {
|
|
server := miniredis.RunT(t)
|
|
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
|
|
oldRedis, oldClient := common.RedisEnabled, common.RDB
|
|
common.RedisEnabled, common.RDB = true, client
|
|
t.Cleanup(func() {
|
|
common.RedisEnabled, common.RDB = oldRedis, oldClient
|
|
liveResults.Delete("group")
|
|
liveResults.Delete("other")
|
|
client.Close()
|
|
})
|
|
ctx := context.Background()
|
|
require.NoError(t, client.Set(ctx, "status:user-result:group", "9999999999999:false", time.Minute).Err())
|
|
assert.Empty(t, LatestUserResults(ctx, []string{"group"}))
|
|
check := func(want UserRequestResult) {
|
|
assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
|
|
liveResults.Delete("group")
|
|
assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
|
|
}
|
|
recordUserResult("other", UserRequestResult{At: 2000, Success: true})
|
|
recordUserResult("group", UserRequestResult{At: 2000, Success: true})
|
|
check(UserRequestResult{At: 2000, Success: true})
|
|
recordUserResult("group", UserRequestResult{At: 2001, Success: false})
|
|
check(UserRequestResult{At: 2001, Success: false})
|
|
recordUserResult("group", UserRequestResult{At: 2000, Success: true})
|
|
check(UserRequestResult{At: 2001, Success: false})
|
|
recordUserResult("group", UserRequestResult{At: 2002, Success: true})
|
|
recordUserResult("group", UserRequestResult{At: 2001, Success: false})
|
|
check(UserRequestResult{At: 2002, Success: true})
|
|
assert.True(t, LatestUserResults(ctx, []string{"other"})["other"].Success)
|
|
ttl, err := client.TTL(ctx, "status:stream-user-result:group").Result()
|
|
require.NoError(t, err)
|
|
assert.Equal(t, 15*time.Minute, ttl)
|
|
server.FastForward(15 * time.Minute)
|
|
assert.Zero(t, LatestUserResults(ctx, []string{"group"})["group"].At)
|
|
}
|
|
|
|
func TestNonStreamRequestsCannotChangeHealth(t *testing.T) {
|
|
oldRedis := common.RedisEnabled
|
|
common.RedisEnabled = false
|
|
t.Cleanup(func() {
|
|
common.RedisEnabled = oldRedis
|
|
liveResults.Delete("stream-only")
|
|
hotBuckets.Clear()
|
|
})
|
|
ctx := context.Background()
|
|
info := &relaycommon.RelayInfo{UsingGroup: "stream-only", OriginModelName: "test-model", TokenId: 1, IsStream: true, StartTime: time.Now()}
|
|
RecordRelayResult(ctx, info, nil)
|
|
want := LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup]
|
|
require.True(t, want.Success)
|
|
require.Positive(t, want.At)
|
|
info.IsStream = false
|
|
for _, apiErr := range []*types.NewAPIError{nil, types.InitOpenAIError("server_error", 500), types.InitOpenAIError("server_error", 502)} {
|
|
RecordRelayResult(ctx, info, apiErr)
|
|
assert.Equal(t, want, LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup])
|
|
}
|
|
info.IsStream = true
|
|
RecordRelayResult(ctx, info, types.InitOpenAIError("server_error", 502))
|
|
assert.False(t, LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup].Success)
|
|
}
|