package api import ( "encoding/json" "fmt" "net/http" "net/http/httptest" "testing" "github.com/alicebob/miniredis/v2" "github.com/ez-api/ez-api/internal/model" "github.com/ez-api/ez-api/internal/service" "github.com/gin-gonic/gin" "github.com/redis/go-redis/v9" "gorm.io/driver/sqlite" "gorm.io/gorm" ) func TestMasterStats_AggregatesByKeyAndModel(t *testing.T) { gin.SetMode(gin.TestMode) dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name()) db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) if err != nil { t.Fatalf("open sqlite: %v", err) } if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.LogRecord{}); err != nil { t.Fatalf("migrate: %v", err) } m := &model.Master{Name: "m1", Group: "g", Status: "active", Epoch: 1, MasterKeyDigest: "d1"} if err := db.Create(m).Error; err != nil { t.Fatalf("create master: %v", err) } k1 := &model.Key{MasterID: m.ID, TokenHash: "h1", Group: "g", Status: "active", IssuedAtEpoch: 1} k2 := &model.Key{MasterID: m.ID, TokenHash: "h2", Group: "g", Status: "active", IssuedAtEpoch: 1} if err := db.Create(k1).Error; err != nil { t.Fatalf("create k1: %v", err) } if err := db.Create(k2).Error; err != nil { t.Fatalf("create k2: %v", err) } if err := db.Create(&model.LogRecord{ Group: "rg", KeyID: k1.ID, ModelName: "ns.m1", ProviderID: 10, ProviderType: "openai", ProviderName: "p1", StatusCode: 200, TokensIn: 5, TokensOut: 7, }).Error; err != nil { t.Fatalf("create log1: %v", err) } if err := db.Create(&model.LogRecord{ Group: "rg", KeyID: k2.ID, ModelName: "ns.m2", ProviderID: 11, ProviderType: "anthropic", ProviderName: "p2", StatusCode: 200, TokensIn: 2, TokensOut: 3, }).Error; err != nil { t.Fatalf("create log2: %v", err) } mr := miniredis.RunT(t) rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) masterSvc := service.NewMasterService(db) syncSvc := service.NewSyncService(rdb) h := NewMasterHandler(db, masterSvc, syncSvc) withMaster := func(next gin.HandlerFunc) gin.HandlerFunc { return func(c *gin.Context) { c.Set("master", m) next(c) } } r := gin.New() r.GET("/v1/stats", withMaster(h.GetSelfStats)) req := httptest.NewRequest(http.MethodGet, "/v1/stats?period=all", nil) rr := httptest.NewRecorder() r.ServeHTTP(rr, req) if rr.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rr.Code, rr.Body.String()) } var resp MasterUsageStatsResponse if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil { t.Fatalf("unmarshal: %v", err) } if resp.TotalRequests != 2 || resp.TotalTokens != 17 { t.Fatalf("unexpected totals: %+v", resp) } if len(resp.ByKey) != 2 || len(resp.ByModel) != 2 { t.Fatalf("unexpected breakdown: %+v", resp) } } func TestAdminStats_AggregatesByProvider(t *testing.T) { gin.SetMode(gin.TestMode) dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", t.Name()) db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) if err != nil { t.Fatalf("open sqlite: %v", err) } if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.LogRecord{}); err != nil { t.Fatalf("migrate: %v", err) } m1 := &model.Master{Name: "m1", Group: "g", Status: "active", Epoch: 1, MasterKeyDigest: "d1"} m2 := &model.Master{Name: "m2", Group: "g", Status: "suspended", Epoch: 1, MasterKeyDigest: "d2"} if err := db.Create(m1).Error; err != nil { t.Fatalf("create m1: %v", err) } if err := db.Create(m2).Error; err != nil { t.Fatalf("create m2: %v", err) } k1 := &model.Key{MasterID: m1.ID, TokenHash: "h1", Group: "g", Status: "active", IssuedAtEpoch: 1} k2 := &model.Key{MasterID: m2.ID, TokenHash: "h2", Group: "g", Status: "active", IssuedAtEpoch: 1} if err := db.Create(k1).Error; err != nil { t.Fatalf("create k1: %v", err) } if err := db.Create(k2).Error; err != nil { t.Fatalf("create k2: %v", err) } if err := db.Create(&model.LogRecord{ Group: "rg", KeyID: k1.ID, ModelName: "ns.m1", ProviderID: 10, ProviderType: "openai", ProviderName: "p1", StatusCode: 200, TokensIn: 4, TokensOut: 6, }).Error; err != nil { t.Fatalf("create log1: %v", err) } if err := db.Create(&model.LogRecord{ Group: "rg", KeyID: k2.ID, ModelName: "ns.m2", ProviderID: 11, ProviderType: "anthropic", ProviderName: "p2", StatusCode: 200, TokensIn: 1, TokensOut: 2, }).Error; err != nil { t.Fatalf("create log2: %v", err) } mr := miniredis.RunT(t) rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()}) masterSvc := service.NewMasterService(db) syncSvc := service.NewSyncService(rdb) adminHandler := NewAdminHandler(db, masterSvc, syncSvc) r := gin.New() r.GET("/admin/stats", adminHandler.GetAdminStats) req := httptest.NewRequest(http.MethodGet, "/admin/stats?period=all", nil) rr := httptest.NewRecorder() r.ServeHTTP(rr, req) if rr.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rr.Code, rr.Body.String()) } var resp AdminUsageStatsResponse if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil { t.Fatalf("unmarshal: %v", err) } if resp.TotalMasters != 2 || resp.ActiveMasters != 1 { t.Fatalf("unexpected master counts: %+v", resp) } if resp.TotalRequests != 2 || resp.TotalTokens != 13 { t.Fatalf("unexpected totals: %+v", resp) } if len(resp.ByProvider) != 2 { t.Fatalf("expected provider breakdown, got %+v", resp.ByProvider) } }