easyai-ai-gateway/apps/api/internal/runner/limits_test.go

30 lines
825 B
Go

package runner
import (
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
)
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"])
}
}