Files
new-api/service/quota_saturation_test.go
T
CaIon bae799ccb1 fix(billing): surface quota saturation events for admin auditing
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.
2026-07-07 12:20:07 +08:00

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")
}