claude-code-proxy/proxy/internal/service/storage_aggregation_test.go

package service

import (
	"testing"

	"github.com/seifghazi/claude-code-monitor/internal/model"
)

func TestAggregationHelpers(t *testing.T) {
	t.Parallel()

	t.Run("addDailyTokens accumulates by date and model", func(t *testing.T) {
		t.Parallel()

		dailyMap := make(map[string]*model.DailyTokens)
		addDailyTokens(dailyMap, "2026-03-20", "claude", 10)
		addDailyTokens(dailyMap, "2026-03-20", "claude", 5)
		addDailyTokens(dailyMap, "2026-03-20", "gpt", 7)

		day := dailyMap["2026-03-20"]
		if day == nil || day.Tokens != 22 || day.Requests != 3 {
			t.Fatalf("unexpected daily aggregate: %#v", day)
		}
		if day.Models["claude"].Tokens != 15 || day.Models["claude"].Requests != 2 {
			t.Fatalf("unexpected claude daily model aggregate: %#v", day.Models["claude"])
		}
		if day.Models["gpt"].Tokens != 7 || day.Models["gpt"].Requests != 1 {
			t.Fatalf("unexpected gpt daily model aggregate: %#v", day.Models["gpt"])
		}
	})

	t.Run("addHourlyTokens accumulates by bucket and model", func(t *testing.T) {
		t.Parallel()

		bucketMap := make(map[string]*model.HourlyTokens)
		addHourlyTokens(bucketMap, "09", "09:00", "claude", 4)
		addHourlyTokens(bucketMap, "09", "09:00", "claude", 6)
		addHourlyTokens(bucketMap, "09", "09:00", "gpt", 2)

		bucket := bucketMap["09"]
		if bucket == nil || bucket.Tokens != 12 || bucket.Requests != 3 || bucket.Label != "09:00" {
			t.Fatalf("unexpected hourly aggregate: %#v", bucket)
		}
		if bucket.Models["claude"].Tokens != 10 || bucket.Models["claude"].Requests != 2 {
			t.Fatalf("unexpected claude hourly model aggregate: %#v", bucket.Models["claude"])
		}
		if bucket.Models["gpt"].Tokens != 2 || bucket.Models["gpt"].Requests != 1 {
			t.Fatalf("unexpected gpt hourly model aggregate: %#v", bucket.Models["gpt"])
		}
	})

	t.Run("addModelTokens accumulates by model", func(t *testing.T) {
		t.Parallel()

		modelMap := make(map[string]*model.ModelTokens)
		addModelTokens(modelMap, "claude", 8)
		addModelTokens(modelMap, "claude", 12)

		got := modelMap["claude"]
		if got == nil || got.Tokens != 20 || got.Requests != 2 {
			t.Fatalf("unexpected model aggregate: %#v", got)
		}
	})
}