new-api/service/channel_test.go

package service

import (
	"errors"
	"net/http"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/setting/operation_setting"
	"github.com/QuantumNous/new-api/types"
	"github.com/stretchr/testify/require"
)

func TestShouldDisableChannelHonorsFeatureFlagAndChannelErrors(t *testing.T) {
	original := common.AutomaticDisableChannelEnabled
	t.Cleanup(func() { common.AutomaticDisableChannelEnabled = original })

	channelErr := types.NewError(errors.New("no usable key"), types.ErrorCodeChannelNoAvailableKey)
	common.AutomaticDisableChannelEnabled = false
	require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, channelErr))

	common.AutomaticDisableChannelEnabled = true
	require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, nil))
	require.True(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, channelErr))

	skipRetry := types.NewError(errors.New("retry later"), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry())
	require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, skipRetry))
}

func TestShouldDisableChannelRecognizesStatusAndOpenAIErrorRules(t *testing.T) {
	originalEnabled := common.AutomaticDisableChannelEnabled
	originalRanges := operation_setting.AutomaticDisableStatusCodeRanges
	t.Cleanup(func() {
		common.AutomaticDisableChannelEnabled = originalEnabled
		operation_setting.AutomaticDisableStatusCodeRanges = originalRanges
	})
	common.AutomaticDisableChannelEnabled = true
	operation_setting.AutomaticDisableStatusCodeRanges = []operation_setting.StatusCodeRange{{Start: http.StatusTooManyRequests, End: http.StatusTooManyRequests}}

	byStatus := types.NewOpenAIError(errors.New("rate limited"), types.ErrorCodeBadResponse, http.StatusTooManyRequests)
	require.True(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, byStatus))

	forbidden := types.NewOpenAIError(errors.New("forbidden"), types.ErrorCodeBadResponse, http.StatusForbidden)
	require.True(t, ShouldDisableChannel(constant.ChannelTypeGemini, forbidden))
	require.False(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, forbidden))

	invalidKey := types.WithOpenAIError(types.OpenAIError{
		Message: "bad key",
		Type:    "invalid_request_error",
		Code:    "invalid_api_key",
	}, http.StatusBadRequest)
	require.True(t, ShouldDisableChannel(constant.ChannelTypeOpenAI, invalidKey))
}

func TestShouldEnableChannelRequiresAutoDisabledStatusWithoutError(t *testing.T) {
	original := common.AutomaticEnableChannelEnabled
	t.Cleanup(func() { common.AutomaticEnableChannelEnabled = original })

	common.AutomaticEnableChannelEnabled = false
	require.False(t, ShouldEnableChannel(nil, common.ChannelStatusAutoDisabled))

	common.AutomaticEnableChannelEnabled = true
	require.False(t, ShouldEnableChannel(types.NewError(errors.New("still failing"), types.ErrorCodeBadResponse), common.ChannelStatusAutoDisabled))
	require.False(t, ShouldEnableChannel(nil, common.ChannelStatusManuallyDisabled))
	require.True(t, ShouldEnableChannel(nil, common.ChannelStatusAutoDisabled))
}