133 lines
3.9 KiB
Go
133 lines
3.9 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
)
|
|
|
|
func TestImageBillingEstimateUsesCountResolutionAndQuality(t *testing.T) {
|
|
service := &Service{}
|
|
candidate := store.RuntimeModelCandidate{
|
|
ModelName: "image-model",
|
|
BaseBillingConfig: map[string]any{
|
|
"image": map[string]any{
|
|
"basePrice": 10,
|
|
"dynamicWeight": map[string]any{
|
|
"resolutionFactors": map[string]any{"2K": 1.5},
|
|
"qualityFactors": map[string]any{"high": 1.5},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
items := service.billings(context.Background(), nil, "images.generations", map[string]any{
|
|
"count": 2,
|
|
"quality": "high",
|
|
"resolution": "2K",
|
|
}, candidate, clients.Response{}, true)
|
|
|
|
line := firstBillingLine(t, items)
|
|
if got, want := floatFromAny(line["amount"]), 45.0; got != want {
|
|
t.Fatalf("image estimated amount = %v, want %v", got, want)
|
|
}
|
|
if got, want := line["quantity"], 2; got != want {
|
|
t.Fatalf("image quantity = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestVideoBillingEstimateUsesFiveSecondUnitsAndDynamicWeights(t *testing.T) {
|
|
service := &Service{}
|
|
candidate := store.RuntimeModelCandidate{
|
|
ModelName: "video-model",
|
|
BaseBillingConfig: map[string]any{
|
|
"video": map[string]any{
|
|
"basePrice": 100,
|
|
"dynamicWeight": map[string]any{
|
|
"resolutionWeights": map[string]any{"1080p": 1.5},
|
|
"audioWeights": map[string]any{"true": 2},
|
|
"referenceVideoWeights": map[string]any{"true": 1.5},
|
|
"voiceSpecifiedWeights": map[string]any{"true": 1.2},
|
|
"unusedCompatibilityField": map[string]any{"true": 99},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
items := service.billings(context.Background(), nil, "videos.generations", map[string]any{
|
|
"audio": true,
|
|
"duration": 12,
|
|
"resolution": "1080p",
|
|
"voice_id": "voice-a",
|
|
"content": []any{
|
|
map[string]any{"type": "video_url", "video_url": map[string]any{"url": "https://example.com/reference.mp4"}},
|
|
},
|
|
}, candidate, clients.Response{}, true)
|
|
|
|
line := firstBillingLine(t, items)
|
|
if got, want := floatFromAny(line["amount"]), 1620.0; got != want {
|
|
t.Fatalf("video estimated amount = %v, want %v", got, want)
|
|
}
|
|
if got, want := floatFromAny(line["durationUnitCount"]), 3.0; got != want {
|
|
t.Fatalf("video duration units = %v, want %v", got, want)
|
|
}
|
|
if got, want := line["quantity"], 3; got != want {
|
|
t.Fatalf("video quantity = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestVideoBillingEstimateSupportsServerMainStyleDynamicKeys(t *testing.T) {
|
|
service := &Service{}
|
|
candidate := store.RuntimeModelCandidate{
|
|
ModelName: "legacy-video-model",
|
|
BaseBillingConfig: map[string]any{
|
|
"videoBase": 100,
|
|
"video": map[string]any{
|
|
"dynamicWeight": map[string]any{
|
|
"720p": 1.25,
|
|
"audio-true": 2,
|
|
"reference-video-true": 1.5,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
items := service.billings(context.Background(), nil, "videos.generations", map[string]any{
|
|
"audio": "true",
|
|
"duration": 5,
|
|
"resolution": "720p",
|
|
"video_list": []any{map[string]any{"url": "https://example.com/reference.mp4"}},
|
|
}, candidate, clients.Response{}, true)
|
|
|
|
line := firstBillingLine(t, items)
|
|
if got, want := floatFromAny(line["amount"]), 375.0; got != want {
|
|
t.Fatalf("legacy video estimated amount = %v, want %v", got, want)
|
|
}
|
|
}
|
|
|
|
func TestVideoDurationEstimateSumsMultiShotDurations(t *testing.T) {
|
|
duration := requestDurationSeconds(map[string]any{
|
|
"multi_prompt": []any{
|
|
map[string]any{"prompt": "shot 1", "duration": 3},
|
|
map[string]any{"prompt": "shot 2", "duration": 7},
|
|
},
|
|
})
|
|
if duration != 10 {
|
|
t.Fatalf("multi-shot duration = %v, want 10", duration)
|
|
}
|
|
}
|
|
|
|
func firstBillingLine(t *testing.T, items []any) map[string]any {
|
|
t.Helper()
|
|
if len(items) != 1 {
|
|
t.Fatalf("items length = %d, want 1: %+v", len(items), items)
|
|
}
|
|
line, ok := items[0].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("item type = %T, want map[string]any", items[0])
|
|
}
|
|
return line
|
|
}
|