Thread int32 saturation clamps from tiered settlement and video task recompute into the consume/task logs under admin_info, so oversized or malformed billing inputs stay auditable. Clamp negative audio duration before token conversion and gate the saturation UI markers on admin.
75 lines
2.6 KiB
Go
75 lines
2.6 KiB
Go
package service
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// TestAttachQuotaSaturationNestsUnderAdminInfo verifies the saturation marker
|
|
// is nested under other.admin_info.quota_saturation so it is admin-only (the
|
|
// log formatter strips admin_info for non-admin viewers).
|
|
func TestAttachQuotaSaturationNestsUnderAdminInfo(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(nil)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
UserId: 7,
|
|
OriginModelName: "gpt-image-1",
|
|
QuotaClamp: &common.QuotaClamp{
|
|
Op: "QuotaFromDecimal",
|
|
Kind: common.QuotaClampOverflow,
|
|
Original: 1.8e19,
|
|
Clamped: common.MaxQuota,
|
|
},
|
|
}
|
|
|
|
other := map[string]interface{}{"model_price": 0.004}
|
|
attachQuotaSaturation(ctx, relayInfo, other)
|
|
|
|
adminInfo, ok := other["admin_info"].(map[string]interface{})
|
|
require.True(t, ok, "admin_info should be created")
|
|
sat, ok := adminInfo["quota_saturation"].(map[string]interface{})
|
|
require.True(t, ok, "quota_saturation should be nested under admin_info")
|
|
require.Equal(t, "QuotaFromDecimal", sat["op"])
|
|
require.Equal(t, common.QuotaClampOverflow, sat["kind"])
|
|
require.Equal(t, common.MaxQuota, sat["clamped"])
|
|
}
|
|
|
|
// TestAttachQuotaSaturationPreservesExistingAdminInfo verifies the marker is
|
|
// merged into a pre-existing admin_info map without clobbering it.
|
|
func TestAttachQuotaSaturationPreservesExistingAdminInfo(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(nil)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{
|
|
QuotaClamp: &common.QuotaClamp{Op: "QuotaFromFloat", Kind: common.QuotaClampUnderflow, Clamped: common.MinQuota},
|
|
}
|
|
other := map[string]interface{}{
|
|
"admin_info": map[string]interface{}{"admin_username": "root"},
|
|
}
|
|
attachQuotaSaturation(ctx, relayInfo, other)
|
|
|
|
adminInfo := other["admin_info"].(map[string]interface{})
|
|
require.Equal(t, "root", adminInfo["admin_username"], "existing admin_info fields preserved")
|
|
require.NotNil(t, adminInfo["quota_saturation"])
|
|
}
|
|
|
|
// TestAttachQuotaSaturationNoClampNoMarker verifies the common case (no
|
|
// saturation) leaves the log untouched.
|
|
func TestAttachQuotaSaturationNoClampNoMarker(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
ctx, _ := gin.CreateTestContext(nil)
|
|
|
|
relayInfo := &relaycommon.RelayInfo{QuotaClamp: nil}
|
|
other := map[string]interface{}{"model_price": 0.004}
|
|
attachQuotaSaturation(ctx, relayInfo, other)
|
|
|
|
_, hasAdmin := other["admin_info"]
|
|
require.False(t, hasAdmin, "no admin_info should be added when there is no clamp")
|
|
}
|