new-api/service/violation_fee_test.go

package service

import (
	"testing"

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

func TestIsViolationFeeCodeRecognizesOnlyReservedPrefix(t *testing.T) {
	require.True(t, IsViolationFeeCode(types.ErrorCode("violation_fee.grok_csam")))
	require.False(t, IsViolationFeeCode(types.ErrorCode("grok_csam")))
}

func TestCalcViolationFeeQuotaRejectsInvalidInputsAndRounds(t *testing.T) {
	oldQuotaPerUnit := common.QuotaPerUnit
	common.QuotaPerUnit = 1000
	t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit })

	require.Equal(t, 0, calcViolationFeeQuota(0, 1))
	require.Equal(t, 0, calcViolationFeeQuota(1, 0))
	require.Equal(t, 333, calcViolationFeeQuota(0.333, 1))
	require.Equal(t, 500, calcViolationFeeQuota(0.25, 2))
}

func TestNormalizeViolationFeeErrorMarksExistingViolationCodeAsNonRetryable(t *testing.T) {
	err := types.WithOpenAIError(types.OpenAIError{
		Code: string(types.ErrorCodeViolationFeeGrokCSAM),
	}, 400)

	normalized := NormalizeViolationFeeError(err)

	require.Equal(t, types.ErrorCodeViolationFeeGrokCSAM, normalized.GetErrorCode())
	require.True(t, types.IsSkipRetryError(normalized))
}