package ai

import "testing"

func TestExtractUsageTokens(t *testing.T) {
	payload := []byte(`{"usage":{"total_tokens":42,"completion_tokens":10}}`)
	if got := extractUsageTokens(payload); got != 42 {
		t.Fatalf("extractUsageTokens() = %d, want 42", got)
	}
}

func TestExtractUsageTokensFallback(t *testing.T) {
	if got := extractUsageTokens([]byte(`{"choices":[]}`)); got != 1 {
		t.Fatalf("expected fallback token count 1, got %d", got)
	}
}