package gateway

import "testing"

func TestUsageCollectorReadsOpenAIJSON(t *testing.T) {
	collector := &usageCollector{}
	collector.feed([]byte(`{"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}}`))
	if got := collector.tokens(); got != 18 {
		t.Fatalf("got %d, want 18", got)
	}
}

func TestUsageCollectorReadsChunkedSSEAndAnthropicUsage(t *testing.T) {
	collector := &usageCollector{sse: true}
	collector.feed([]byte("event: message_start\ndata: {\"message\":{\"usage\":{\"input_tokens\":13}}}\n\n"))
	collector.feed([]byte("event: message_delta\ndata: {\"usage\":{\"output_"))
	collector.feed([]byte("tokens\":9}}\n\ndata: [DONE]\n\n"))
	if got := collector.tokens(); got != 22 {
		t.Fatalf("got %d, want 22", got)
	}
}