new-api/service/usage_helpr_test.go

package service

import (
	"net/http/httptest"
	"testing"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"github.com/QuantumNous/new-api/dto"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func TestResponseText2UsageSetsLocalCountFlagAndTotals(t *testing.T) {
	ctx, _ := gin.CreateTestContext(httptest.NewRecorder())

	usage := ResponseText2Usage(ctx, "hello", "gpt-5", 3)

	require.Equal(t, 3, usage.PromptTokens)
	require.Equal(t, EstimateTokenByModel("gpt-5", "hello"), usage.CompletionTokens)
	require.Equal(t, usage.PromptTokens+usage.CompletionTokens, usage.TotalTokens)
	require.True(t, common.GetContextKeyBool(ctx, constant.ContextKeyLocalCountTokens))
}

func TestValidUsageRequiresAtLeastOneNonZeroTokenCount(t *testing.T) {
	require.False(t, ValidUsage(nil))
	require.False(t, ValidUsage(&dto.Usage{}))
	require.True(t, ValidUsage(&dto.Usage{PromptTokens: 1}))
	require.True(t, ValidUsage(&dto.Usage{CompletionTokens: 1}))
}