package perfmetrics

import (
	"context"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func TestLiveUserSuccess(t *testing.T) {
	t.Cleanup(func() {
		for _, name := range []string{"user", "probe", "internal", "failed"} {
			liveSuccess.Delete(name)
		}
	})
	for _, tc := range []struct {
		name  string
		test  bool
		token int
		fail  bool
		want  bool
	}{
		{"user", false, 1, false, true}, {"probe", true, 1, false, false}, {"internal", false, 0, false, false}, {"failed", false, 1, true, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var apiErr *types.NewAPIError
			if tc.fail {
				apiErr = types.InitOpenAIError("server_error", 500)
			}
			RecordRelayResult(context.Background(), &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, apiErr)
			assert.Equal(t, tc.want, LatestUserSuccess(context.Background(), []string{tc.name})[tc.name] > 0)
		})
	}
}