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