gpt-load/internal/control/usage_rolling_test.go

package control

import (
	"encoding/json"
	"fmt"
	"net/http"
	"testing"
	"time"

	"github.com/gin-gonic/gin"

	"gpt-load/internal/platform/config"
	"gpt-load/internal/requestlog"
	"gpt-load/internal/storage/models"
)

func TestUsageAPIRollingHourReadsLogsAndScopesAccessKey(t *testing.T) {
	t.Parallel()
	initControlI18n(t)
	fixture := newServiceFixture(t)
	key, err := fixture.service.CreateAccessKey(t.Context(), AccessKeyCreateRequest{Name: "rolling usage viewer"})
	if err != nil {
		t.Fatalf("create usage access key: %v", err)
	}
	now := time.Date(2026, time.September, 7, 14, 5, 23, 123_000_000, time.UTC)
	from := now.Add(-time.Hour)
	for index, item := range []struct {
		completedAt time.Time
		accessKeyID uint
		attempts    int
	}{
		{from.Add(-time.Millisecond), key.ID, 1},
		{from, key.ID, 1},
		{now.Add(-time.Millisecond), key.ID, 1},
		{now, key.ID, 1},
		{now.Add(-time.Minute), key.ID, 0},
		{now.Add(-time.Second), key.ID + 1, 1},
	} {
		row := models.RequestLog{
			ID:            fmt.Sprintf("00000000-0000-4000-8000-%012d", index+1),
			CompletedAtMS: item.completedAt.UnixMilli(), AccessKeyID: item.accessKeyID,
			GroupID: 7, ChannelID: "openai", CredentialID: 11,
			Protocol: "openai-completions", ClientModel: "usage-model", UpstreamModel: "usage-model",
			Status: "success", StatusCode: 200, AttemptCount: item.attempts,
			UsageState: "complete", CostState: "priced", PricingCompleteness: "complete",
			UncachedInputTokens: 3, OutputTokens: 5, EstimatedCostNanoUSD: 700,
		}
		if err := fixture.db.Create(&row).Error; err != nil {
			t.Fatalf("create usage request log: %v", err)
		}
	}
	if err := fixture.db.Create(&models.UsageStat{
		BucketStartMS: from.Truncate(time.Hour).UnixMilli(), AccessKeyID: key.ID,
		GroupID: 7, ChannelID: "openai", CredentialID: 11, Model: "usage-model",
		RequestCount: 99, SuccessCount: 99,
	}).Error; err != nil {
		t.Fatalf("create independent hourly usage stat: %v", err)
	}
	fixture.service.now = func() time.Time { return now }
	fixture.service.usageStats = requestlog.NewService(fixture.db, nil, nil)
	engine := gin.New()
	NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
	for _, test := range []struct {
		name      string
		auth      string
		query     string
		wantCount int64
		wantCost  string
		wantGrain requestlog.UsageGranularity
	}{
		{"management frozen hour", "test-auth-key", usageTestTimeQuery(from.Add(-time.Minute), now.Add(-time.Minute)), 2, "1400", requestlog.UsageGranularityMinute},
		{"management rolling hour", "test-auth-key", usageTestTimeQuery(from, now), 3, "2100", requestlog.UsageGranularityMinute},
		{"access key rolling hour", key.Key, usageTestTimeQuery(from, now), 2, "1400", requestlog.UsageGranularityMinute},
		{"management filters access key", "test-auth-key", usageTestTimeQuery(from, now) + fmt.Sprintf("&access_key_id=%d&group_id=7", key.ID), 2, "1400", requestlog.UsageGranularityMinute},
		{"daily range combines hourly and boundary sources", "test-auth-key", usageTestTimeQuery(now.Add(-24*time.Hour), now), 101, "1400", requestlog.UsageGranularityHour},
		{"aligned hour uses logs and minute buckets", "test-auth-key",
			fmt.Sprintf("from_ms=%d&to_ms=%d", from.Truncate(time.Hour).UnixMilli(), from.Truncate(time.Hour).Add(time.Hour).UnixMilli()),
			2, "1400", requestlog.UsageGranularityMinute},
	} {
		t.Run(test.name, func(t *testing.T) {
			recorder := performUsageRequest(engine, test.auth, test.query)
			if recorder.Code != http.StatusOK {
				t.Fatalf("usage response = %d %s", recorder.Code, recorder.Body.String())
			}
			var envelope struct {
				Data usageResponse `json:"data"`
			}
			if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
				t.Fatalf("decode usage response: %v", err)
			}
			data := envelope.Data
			if data.Granularity != test.wantGrain {
				t.Fatalf("usage granularity = %s, want %s", data.Granularity, test.wantGrain)
			}
			if data.Summary.RequestCount != test.wantCount || data.Summary.EstimatedCostNanoUSD != test.wantCost {
				t.Fatalf("usage summary = %+v", data.Summary)
			}
			if test.query == usageTestTimeQuery(from, now) {
				if len(data.Series) != 2 || data.Series[0].BucketStartMS != from.UnixMilli() ||
					data.Series[1].BucketEndMS != now.UnixMilli() || data.Granularity != "minute" {
					t.Fatalf("rolling usage series = %+v", data.Series)
				}
				if data.Series[0].RequestCount+data.Series[1].RequestCount != test.wantCount {
					t.Fatalf("series differs from summary: %+v", data.Series)
				}
				for _, distribution := range data.Distributions.Model {
					if len(distribution.Items) != 1 || distribution.Items[0].RequestCount != test.wantCount ||
						distribution.Items[0].EstimatedCostNanoUSD != test.wantCost {
						t.Fatalf("model distribution differs from summary: %+v", distribution)
					}
				}
			}
			if test.auth == "test-auth-key" && (data.Distributions.Group == nil || data.CollectionHealth.Scope != "current_process") {
				t.Fatalf("filtered management response lost management view: %+v", data)
			}
			if test.auth == key.Key && (data.Distributions.Group != nil || data.Distributions.AccessKey != nil ||
				data.CollectionHealth.Scope != "access_key") {
				t.Fatalf("access key response exposes management scope: %+v", data)
			}
		})
	}
	for _, id := range []uint{key.ID, key.ID + 1} {
		recorder := performUsageRequest(engine, key.Key, usageTestTimeQuery(from, now)+fmt.Sprintf("&access_key_id=%d", id))
		if recorder.Code != http.StatusBadRequest {
			t.Fatalf("access key cannot select a scope: %d %s", recorder.Code, recorder.Body.String())
		}
	}
}