mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-09-08 19:06:22 +00:00
70 lines
2.8 KiB
Go
70 lines
2.8 KiB
Go
package model
|
|
|
|
import (
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func pricingUsagePluginSource(version, usageSchema string) string {
|
|
return fmt.Sprintf(`
|
|
export const meta = {
|
|
apiVersion: 1, key: "pricing-usage-probe", name: "Pricing Usage Probe", version: %q, author: {name: "Test"},
|
|
models: ["pricing-usage-model"], fetchMode: "per_task", usageSchema: %s
|
|
};
|
|
export function buildSubmitRequest() { return {}; }
|
|
export function parseSubmitResponse() { return {}; }
|
|
export function buildQueryRequest() { return {}; }
|
|
export function parseTaskResult() { return {}; }
|
|
`, version, usageSchema)
|
|
}
|
|
|
|
func TestPricingCarriesTaskUsageSchemaAndRefreshesWithPluginGeneration(t *testing.T) {
|
|
resetPricingEndpointTestTables(t)
|
|
const pluginKey = "pricing-usage-probe"
|
|
initialSource := pricingUsagePluginSource("1.0.0", `{
|
|
seconds: {type: "number", unit: "second", description: "Estimated duration."}
|
|
}`)
|
|
_, err := jsplugin.DefaultRegistry.Register(initialSource, jsplugin.Options{})
|
|
require.NoError(t, err)
|
|
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) })
|
|
|
|
insertPricingEndpointChannel(t, 901, constant.ChannelTypeTaskPlugin, dto.ChannelOtherSettings{})
|
|
insertPricingEndpointAbility(t, 901, "pricing-usage-model")
|
|
insertPricingEndpointAbility(t, 901, "ordinary-model")
|
|
|
|
initialPricing := pricingByModel(GetPricing())
|
|
require.Contains(t, initialPricing, "pricing-usage-model")
|
|
require.Contains(t, initialPricing, "ordinary-model")
|
|
assert.Equal(t, "second", initialPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Unit)
|
|
assert.Equal(t, "Estimated duration.", initialPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Description["en"])
|
|
assert.Nil(t, initialPricing["ordinary-model"].BillingUsageSchema)
|
|
|
|
updatedSource := pricingUsagePluginSource("1.1.0", `{
|
|
seconds: {type: "number", unit: "second", description: "Measured duration."},
|
|
clips: {type: "number", unit: "count", description: "Generated clip count."}
|
|
}`)
|
|
_, err = jsplugin.DefaultRegistry.Register(updatedSource, jsplugin.Options{})
|
|
require.NoError(t, err)
|
|
lastGetPricingTime = time.Now().Add(-2 * time.Minute)
|
|
|
|
refreshedPricing := pricingByModel(GetPricing())
|
|
require.Len(t, refreshedPricing["pricing-usage-model"].BillingUsageSchema, 2)
|
|
assert.Equal(t, "Measured duration.", refreshedPricing["pricing-usage-model"].BillingUsageSchema["seconds"].Description["en"])
|
|
assert.Equal(t, "count", refreshedPricing["pricing-usage-model"].BillingUsageSchema["clips"].Unit)
|
|
}
|
|
|
|
func pricingByModel(pricings []Pricing) map[string]Pricing {
|
|
result := make(map[string]Pricing, len(pricings))
|
|
for _, pricing := range pricings {
|
|
result[pricing.ModelName] = pricing
|
|
}
|
|
return result
|
|
}
|