135 lines
4.2 KiB
Go
135 lines
4.2 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 TestTokenUsageAmountsUsesActualUsageForTPM(t *testing.T) {
|
|
got := tokenUsageAmounts(clients.Usage{InputTokens: 12, OutputTokens: 8, TotalTokens: 21})
|
|
|
|
if got["tpm_input"] != 12 {
|
|
t.Fatalf("expected input token amount 12, got %v", got["tpm_input"])
|
|
}
|
|
if got["tpm_output"] != 8 {
|
|
t.Fatalf("expected output token amount 8, got %v", got["tpm_output"])
|
|
}
|
|
if got["tpm_total"] != 21 {
|
|
t.Fatalf("expected total token amount 21, got %v", got["tpm_total"])
|
|
}
|
|
}
|
|
|
|
func TestTokenUsageAmountsFallsBackToInputOutputTotal(t *testing.T) {
|
|
got := tokenUsageAmounts(clients.Usage{InputTokens: 3, OutputTokens: 5})
|
|
|
|
if got["tpm_total"] != 8 {
|
|
t.Fatalf("expected total token fallback 8, got %v", got["tpm_total"])
|
|
}
|
|
}
|
|
|
|
func TestBillingsSplitsCachedInputTokens(t *testing.T) {
|
|
service := &Service{}
|
|
candidate := store.RuntimeModelCandidate{
|
|
ModelName: "test-model",
|
|
BillingConfig: map[string]any{
|
|
"textInputPer1k": 1.0,
|
|
"textCachedInputPer1k": 0.25,
|
|
"textOutputPer1k": 2.0,
|
|
},
|
|
}
|
|
response := clients.Response{
|
|
Usage: clients.Usage{
|
|
InputTokens: 1000,
|
|
CachedInputTokens: 400,
|
|
OutputTokens: 500,
|
|
TotalTokens: 1500,
|
|
},
|
|
}
|
|
|
|
lines := service.billings(context.Background(), nil, "chat.completions", nil, candidate, response, false)
|
|
|
|
if len(lines) != 3 {
|
|
t.Fatalf("expected uncached input, cached input, and output lines, got %+v", lines)
|
|
}
|
|
uncached, _ := lines[0].(map[string]any)
|
|
cached, _ := lines[1].(map[string]any)
|
|
output, _ := lines[2].(map[string]any)
|
|
if uncached["resourceType"] != "text_input" || uncached["quantity"] != 600 || uncached["amount"] != 0.6 {
|
|
t.Fatalf("unexpected uncached input billing line: %+v", uncached)
|
|
}
|
|
if cached["resourceType"] != "text_cached_input" || cached["quantity"] != 400 || cached["amount"] != 0.1 {
|
|
t.Fatalf("unexpected cached input billing line: %+v", cached)
|
|
}
|
|
if output["resourceType"] != "text_output" || output["quantity"] != 500 || output["amount"] != 1.0 {
|
|
t.Fatalf("unexpected output billing line: %+v", output)
|
|
}
|
|
}
|
|
|
|
func TestBillingsDefaultsCachedInputPriceToOneTenthOfInputPrice(t *testing.T) {
|
|
service := &Service{}
|
|
candidate := store.RuntimeModelCandidate{
|
|
ModelName: "test-model",
|
|
BillingConfig: map[string]any{
|
|
"textInputPer1k": 1.0,
|
|
"textOutputPer1k": 2.0,
|
|
},
|
|
}
|
|
response := clients.Response{
|
|
Usage: clients.Usage{
|
|
InputTokens: 1000,
|
|
CachedInputTokens: 400,
|
|
OutputTokens: 500,
|
|
TotalTokens: 1500,
|
|
},
|
|
}
|
|
|
|
lines := service.billings(context.Background(), nil, "chat.completions", nil, candidate, response, false)
|
|
|
|
if len(lines) != 3 {
|
|
t.Fatalf("expected uncached input, cached input, and output lines, got %+v", lines)
|
|
}
|
|
cached, _ := lines[1].(map[string]any)
|
|
if cached["resourceType"] != "text_cached_input" || cached["quantity"] != 400 || cached["amount"] != 0.04 {
|
|
t.Fatalf("unexpected cached input fallback billing line: %+v", cached)
|
|
}
|
|
}
|
|
|
|
func TestEffectiveRateLimitPolicyTreatsEmptyRuntimePolicyAsUnlimited(t *testing.T) {
|
|
policy := effectiveRateLimitPolicy(store.RuntimeModelCandidate{
|
|
PlatformRateLimitPolicy: map[string]any{"rules": []any{
|
|
map[string]any{"metric": "rpm", "limit": 500},
|
|
}},
|
|
RuntimePolicySetID: "runtime-policy-1",
|
|
RuntimeRateLimitPolicy: map[string]any{"rules": []any{}},
|
|
})
|
|
|
|
if hasRules(policy) {
|
|
t.Fatalf("expected empty runtime policy to clear inherited limits, got %+v", policy)
|
|
}
|
|
}
|
|
|
|
func TestReservationsFromPolicySkipsNonPositiveLimits(t *testing.T) {
|
|
reservations := reservationsFromPolicy(
|
|
"platform_model",
|
|
"model-1",
|
|
"Model 1",
|
|
nil,
|
|
map[string]any{"rules": []any{
|
|
map[string]any{"metric": "rpm", "limit": -1},
|
|
map[string]any{"metric": "tpm_total", "limit": 0},
|
|
map[string]any{"metric": "concurrent", "limit": 2},
|
|
}},
|
|
map[string]any{"prompt": "hello"},
|
|
)
|
|
|
|
if len(reservations) != 1 {
|
|
t.Fatalf("expected only the positive concurrent rule to reserve, got %+v", reservations)
|
|
}
|
|
if reservations[0].Metric != "concurrent" || reservations[0].Limit != 2 {
|
|
t.Fatalf("expected concurrent reservation with limit 2, got %+v", reservations[0])
|
|
}
|
|
}
|