0220df8429
Pre-fill BillingRequestInput from dto.Request before ModelPriceHelper, so tiered_expr billing resolves param() from the structured request instead of reading HTTP body (which is empty in channel-test context). - attachTestBillingRequestInput: marshal dto.Request → RequestInput - ResolveIncomingBillingExprRequestInput: early-return when pre-filled - settleTestQuota / buildTestLogOther: align test settlement & logging with production TryTieredSettle / InjectTieredBillingInfo paths
64 lines
1.9 KiB
Go
64 lines
1.9 KiB
Go
package helper
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/samber/lo"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func TestResolveIncomingBillingExprRequestInput(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
ctx, _ := gin.CreateTestContext(recorder)
|
|
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
|
ctx.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
body := []byte(`{"service_tier":"fast"}`)
|
|
ctx.Request.Body = io.NopCloser(bytes.NewReader(body))
|
|
ctx.Set(common.KeyRequestBody, body)
|
|
|
|
info := &relaycommon.RelayInfo{
|
|
RequestHeaders: map[string]string{"Content-Type": "application/json"},
|
|
}
|
|
|
|
input, err := ResolveIncomingBillingExprRequestInput(ctx, info)
|
|
require.NoError(t, err)
|
|
require.Equal(t, body, input.Body)
|
|
require.Equal(t, "application/json", input.Headers["Content-Type"])
|
|
}
|
|
|
|
func TestBuildBillingExprRequestInputFromRequest(t *testing.T) {
|
|
request := &dto.GeneralOpenAIRequest{
|
|
Model: "gemini-3.1-pro-preview",
|
|
Stream: lo.ToPtr(true),
|
|
Messages: []dto.Message{
|
|
{
|
|
Role: "user",
|
|
Content: "hi",
|
|
},
|
|
},
|
|
MaxTokens: lo.ToPtr(uint(3000)),
|
|
}
|
|
|
|
input, err := BuildBillingExprRequestInputFromRequest(request, map[string]string{
|
|
"Content-Type": "application/json",
|
|
"X-Test": "1",
|
|
})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "application/json", input.Headers["Content-Type"])
|
|
require.Equal(t, "1", input.Headers["X-Test"])
|
|
require.True(t, gjson.GetBytes(input.Body, "stream").Bool())
|
|
require.Equal(t, "user", gjson.GetBytes(input.Body, "messages.0.role").String())
|
|
require.Equal(t, float64(3000), gjson.GetBytes(input.Body, "max_tokens").Float())
|
|
}
|