Files
ez-api/internal/api/auth_handler_test.go
RC-CHN 1ee6bea413 feat(api): enhance whoami endpoint with realtime stats and extended key info
Add realtime statistics (requests, tokens, QPS, rate limiting) to whoami
response for both master and key authentication types. Extend key response
with additional fields including master name, model limits, quota tracking,
and usage statistics.

- Inject StatsService into AuthHandler for realtime stats retrieval
- Add WhoamiRealtimeView struct for realtime statistics
- Include admin permissions field in admin response
- Add comprehensive key metadata (quotas, model limits, usage stats)
- Add test for expired key returning 401 Unauthorized
2026-01-06 09:15:49 +08:00

209 lines
6.1 KiB
Go

package api
import (
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/alicebob/miniredis/v2"
"github.com/ez-api/ez-api/internal/model"
"github.com/ez-api/ez-api/internal/service"
"github.com/ez-api/foundation/tokenhash"
"github.com/gin-gonic/gin"
"github.com/redis/go-redis/v9"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
)
func newAuthHandler(t *testing.T) (*AuthHandler, *gorm.DB, *miniredis.Miniredis) {
t.Helper()
gin.SetMode(gin.TestMode)
t.Setenv("EZ_ADMIN_TOKEN", "admin-secret")
adminService, err := service.NewAdminService()
if err != nil {
t.Fatalf("NewAdminService: %v", err)
}
db, err := gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
if err != nil {
t.Fatalf("open sqlite: %v", err)
}
if err := db.AutoMigrate(&model.Master{}, &model.Key{}); err != nil {
t.Fatalf("migrate: %v", err)
}
mr := miniredis.RunT(t)
rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
masterService := service.NewMasterService(db)
statsService := service.NewStatsService(rdb)
handler := NewAuthHandler(db, rdb, adminService, masterService, statsService)
return handler, db, mr
}
func TestAuthHandler_Whoami_InvalidIssuedAtEpoch_Returns401(t *testing.T) {
handler, _, mr := newAuthHandler(t)
token := "sk-live-invalid-epoch"
hash := tokenhash.HashToken(token)
mr.HSet("auth:token:"+hash, "master_id", "1")
mr.HSet("auth:token:"+hash, "issued_at_epoch", "bad")
mr.HSet("auth:token:"+hash, "status", "active")
mr.HSet("auth:token:"+hash, "group", "default")
mr.HSet("auth:master:1", "epoch", "1")
mr.HSet("auth:master:1", "status", "active")
r := gin.New()
r.GET("/auth/whoami", handler.Whoami)
req := httptest.NewRequest(http.MethodGet, "/auth/whoami", nil)
req.Header.Set("Authorization", "Bearer "+token)
rr := httptest.NewRecorder()
r.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d body=%s", rr.Code, rr.Body.String())
}
}
func TestAuthHandler_Whoami_InvalidMasterEpoch_Returns401(t *testing.T) {
handler, _, mr := newAuthHandler(t)
token := "sk-live-invalid-master-epoch"
hash := tokenhash.HashToken(token)
mr.HSet("auth:token:"+hash, "master_id", "1")
mr.HSet("auth:token:"+hash, "issued_at_epoch", "1")
mr.HSet("auth:token:"+hash, "status", "active")
mr.HSet("auth:token:"+hash, "group", "default")
mr.HSet("auth:master:1", "epoch", "bad")
mr.HSet("auth:master:1", "status", "active")
r := gin.New()
r.GET("/auth/whoami", handler.Whoami)
req := httptest.NewRequest(http.MethodGet, "/auth/whoami", nil)
req.Header.Set("Authorization", "Bearer "+token)
rr := httptest.NewRecorder()
r.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected 401, got %d body=%s", rr.Code, rr.Body.String())
}
}
func TestAuthHandler_Whoami_KeyResponseIncludesIPRules(t *testing.T) {
handler, db, mr := newAuthHandler(t)
token := "sk-live-valid"
hash := tokenhash.HashToken(token)
expiresAt := time.Now().Add(time.Hour).Unix()
mr.HSet("auth:token:"+hash, "master_id", "1")
mr.HSet("auth:token:"+hash, "issued_at_epoch", "1")
mr.HSet("auth:token:"+hash, "status", "active")
mr.HSet("auth:token:"+hash, "group", "default")
mr.HSet("auth:token:"+hash, "expires_at", fmt.Sprintf("%d", expiresAt))
mr.HSet("auth:master:1", "epoch", "1")
mr.HSet("auth:master:1", "status", "active")
key := &model.Key{
MasterID: 1,
TokenHash: hash,
Group: "default",
Scopes: "chat:write",
DefaultNamespace: "default",
Namespaces: "default",
Status: "active",
IssuedAtEpoch: 1,
IssuedBy: "master",
AllowIPs: "1.2.3.4",
DenyIPs: "5.6.7.0/24",
}
if err := db.Create(key).Error; err != nil {
t.Fatalf("create key: %v", err)
}
r := gin.New()
r.GET("/auth/whoami", handler.Whoami)
req := httptest.NewRequest(http.MethodGet, "/auth/whoami", nil)
req.Header.Set("Authorization", "Bearer "+token)
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 map[string]any
if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
t.Fatalf("decode response: %v", err)
}
if resp["allow_ips"] != "1.2.3.4" {
t.Fatalf("expected allow_ips, got %v", resp["allow_ips"])
}
if resp["deny_ips"] != "5.6.7.0/24" {
t.Fatalf("expected deny_ips, got %v", resp["deny_ips"])
}
if got, ok := resp["expires_at"].(float64); !ok || int64(got) != expiresAt {
t.Fatalf("expected expires_at=%d, got %v", expiresAt, resp["expires_at"])
}
}
func TestAuthHandler_Whoami_ExpiredKey_Returns401(t *testing.T) {
handler, db, mr := newAuthHandler(t)
token := "sk-live-expired"
hash := tokenhash.HashToken(token)
// Set expiration to 1 hour ago
expiresAt := time.Now().Add(-time.Hour).Unix()
mr.HSet("auth:token:"+hash, "master_id", "1")
mr.HSet("auth:token:"+hash, "issued_at_epoch", "1")
mr.HSet("auth:token:"+hash, "status", "active")
mr.HSet("auth:token:"+hash, "group", "default")
mr.HSet("auth:token:"+hash, "expires_at", fmt.Sprintf("%d", expiresAt))
mr.HSet("auth:master:1", "epoch", "1")
mr.HSet("auth:master:1", "status", "active")
// Create key in DB
key := &model.Key{
MasterID: 1,
TokenHash: hash,
Group: "default",
Scopes: "chat:write",
DefaultNamespace: "default",
Namespaces: "default",
Status: "active",
IssuedAtEpoch: 1,
IssuedBy: "master",
}
if err := db.Create(key).Error; err != nil {
t.Fatalf("create key: %v", err)
}
r := gin.New()
r.GET("/auth/whoami", handler.Whoami)
req := httptest.NewRequest(http.MethodGet, "/auth/whoami", nil)
req.Header.Set("Authorization", "Bearer "+token)
rr := httptest.NewRecorder()
r.ServeHTTP(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("expected 401 for expired key, got %d body=%s", rr.Code, rr.Body.String())
}
var resp map[string]any
if err := json.Unmarshal(rr.Body.Bytes(), &resp); err != nil {
t.Fatalf("decode response: %v", err)
}
if resp["error"] != "token has expired" {
t.Fatalf("expected 'token has expired' error, got %v", resp["error"])
}
}