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