new-api/setting/ratio_setting/model_display_pricing_test.go

package ratio_setting

import (
	"testing"

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

func TestModelDisplayPricingRoundTripNormalizeAndSort(t *testing.T) {
	require.NoError(t, UpdateModelDisplayPricingByJSONString(`{}`))

	jsonStr := `{
		"kling-video": [
			{"specification":"1080p","official_supplier_tip":"1080P","price":0.58333,"unit":"second","sort_order":3},
			{"specification":"768p-6s","official_supplier_tip":"768P 6s","price":0.33333,"unit":"second","sort_order":1,"discount_rate":0.8},
			{"specification":"768p-10s","official_supplier_tip":"768P 10s","price":0.4,"unit":"second","sort_order":2,"discount_rate":0}
		]
	}`

	require.NoError(t, UpdateModelDisplayPricingByJSONString(jsonStr))
	items := GetModelDisplayPricing("kling-video")
	require.Len(t, items, 3)
	require.Equal(t, "768p-6s", items[0].Specification)
	require.Equal(t, 0.8, items[0].DiscountRate)
	require.Equal(t, "768p-10s", items[1].Specification)
	require.Equal(t, 1.0, items[1].DiscountRate)
	require.Equal(t, "1080p", items[2].Specification)
	require.Equal(t, 1.0, items[2].DiscountRate)

	items[0].Specification = "mutated"
	require.Equal(t, "768p-6s", GetModelDisplayPricing("kling-video")[0].Specification)

	roundTrip := ModelDisplayPricing2JSONString()
	require.NoError(t, UpdateModelDisplayPricingByJSONString(roundTrip))
	require.Equal(t, "768p-6s", GetModelDisplayPricing("kling-video")[0].Specification)
}

func TestSetModelDisplayPricingValidation(t *testing.T) {
	tests := []struct {
		name  string
		model string
		items []types.ModelDisplayPricingItem
		want  string
	}{
		{
			name:  "empty_model",
			model: "",
			items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: 1, Unit: "second"}},
			want:  "model name is required",
		},
		{
			name:  "empty_specification",
			model: "kling-video",
			items: []types.ModelDisplayPricingItem{{Price: 1, Unit: "second"}},
			want:  "specification is required",
		},
		{
			name:  "negative_price",
			model: "kling-video",
			items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: -1, Unit: "second"}},
			want:  "price must be >= 0",
		},
		{
			name:  "empty_unit",
			model: "kling-video",
			items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: 1}},
			want:  "unit is required",
		},
		{
			name:  "bad_discount",
			model: "kling-video",
			items: []types.ModelDisplayPricingItem{{Specification: "768p", Price: 1, Unit: "second", DiscountRate: 1.2}},
			want:  "discount_rate must be > 0 and <= 1",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			err := SetModelDisplayPricing(tt.model, tt.items)
			require.Error(t, err)
			require.Contains(t, err.Error(), tt.want)
		})
	}
}

func TestDeleteModelDisplayPricing(t *testing.T) {
	require.NoError(t, UpdateModelDisplayPricingByJSONString(`{}`))
	require.NoError(t, SetModelDisplayPricing("kling-video", []types.ModelDisplayPricingItem{
		{Specification: "768p", Price: 1, Unit: "second"},
	}))
	require.Len(t, GetModelDisplayPricing("kling-video"), 1)
	require.NoError(t, DeleteModelDisplayPricing("kling-video"))
	require.Empty(t, GetModelDisplayPricing("kling-video"))
}