Files
NianAIGC/backend/internal/httpapi/usage_test.go

95 lines
4.7 KiB
Go

package httpapi
import (
"context"
"errors"
"fmt"
"net/http"
"testing"
"time"
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/identity"
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/usage"
)
func TestUsageRoutesEnforcePersonalAndAdminScope(t *testing.T) {
reporter := &usageReporterStub{}
user := usageTestHandler(t, identity.Session{AuthMode: identity.AuthModeUser, User: identity.User{ID: "db-user", ClientID: "platform", OrganizationID: "db-org", Role: "user"}}, reporter)
if got := serveJSON(t, user, http.MethodGet, "/api/usage?preset=7d", nil); got.Code != 200 || reporter.personal.AccountID != "db-user" || reporter.personal.Preset != usage.Preset7Days {
t.Fatalf("status=%d request=%+v", got.Code, reporter.personal)
}
if got := serveJSON(t, user, http.MethodGet, "/api/admin/usage", nil); got.Code != 403 {
t.Fatalf("user admin status=%d", got.Code)
}
orgAdmin := usageTestHandler(t, identity.Session{AuthMode: identity.AuthModeAdmin, User: identity.User{ID: "admin", ClientID: "platform", OrganizationID: "db-org", Role: "organization_admin"}}, reporter)
got := serveJSON(t, orgAdmin, http.MethodGet, "/api/admin/usage?organizationId=attacker&ownerId=other&provider=bailian", nil)
if got.Code != 200 || reporter.admin.OrganizationID != "db-org" || reporter.admin.OwnerID != "" || reporter.admin.Provider != "bailian" || !reporter.admin.RedactAccounts {
t.Fatalf("status=%d request=%+v body=%s", got.Code, reporter.admin, got.Body.String())
}
super := usageTestHandler(t, identity.Session{AuthMode: identity.AuthModeAdmin, User: identity.User{ID: "root", ClientID: "platform", Role: "super_admin"}}, reporter)
got = serveJSON(t, super, http.MethodGet, "/api/admin/usage?organizationId=org-x&ownerId=user-x&capability=image.generate", nil)
if got.Code != 200 || reporter.admin.OrganizationID != "org-x" || reporter.admin.OwnerID != "user-x" || reporter.admin.RedactAccounts {
t.Fatalf("status=%d request=%+v", got.Code, reporter.admin)
}
}
func TestUsageValidatesFiltersAndHidesInfrastructureErrors(t *testing.T) {
reporter := &usageReporterStub{}
h := usageTestHandler(t, identity.Session{AuthMode: identity.AuthModeAdmin, User: identity.User{ID: "root", ClientID: "platform", Role: "super_admin"}}, reporter)
for _, path := range []string{"/api/admin/usage?capability=bad", "/api/admin/usage?provider=mock", "/api/admin/usage?startDate=not-date"} {
if got := serveJSON(t, h, http.MethodGet, path, nil); got.Code != 400 {
t.Fatalf("%s status=%d", path, got.Code)
}
}
reporter.err = context.DeadlineExceeded
if got := serveJSON(t, h, http.MethodGet, "/api/admin/usage", nil); got.Code != 500 || got.Body.String() != "{\"error\":\"服务器内部错误。\"}\n" {
t.Fatalf("status/body=%d %q", got.Code, got.Body.String())
}
}
func TestUsageMapsDomainDateRangeErrorsToBadRequest(t *testing.T) {
reporter := &usageReporterStub{}
h := usageTestHandler(t, identity.Session{AuthMode: identity.AuthModeAdmin, User: identity.User{ID: "root", ClientID: "platform", Role: "super_admin"}}, reporter)
reporter.err = usage.ErrDateRangeTooLong
if got := serveJSON(t, h, http.MethodGet, "/api/admin/usage?startDate=2010-01-01&endDate=2026-01-01", nil); got.Code != 400 || got.Body.String() != "{\"error\":\"单次查询最多支持 10 年。\"}\n" {
t.Fatalf("long range status/body=%d %q", got.Code, got.Body.String())
}
reporter.err = fmt.Errorf("wrapped: %w", usage.ErrInvalidDateRange)
if got := serveJSON(t, h, http.MethodGet, "/api/admin/usage?startDate=2026-08-13&endDate=2026-08-12", nil); got.Code != 400 || got.Body.String() != "{\"error\":\"开始日期不能晚于结束日期。\"}\n" {
t.Fatalf("reverse range status/body=%d %q", got.Code, got.Body.String())
}
reporter.err = errors.New("database unavailable")
if got := serveJSON(t, h, http.MethodGet, "/api/admin/usage", nil); got.Code != 500 {
t.Fatalf("infrastructure status=%d", got.Code)
}
}
type usageReporterStub struct {
personal usage.PersonalRequest
admin usage.AdminRequest
err error
}
func (s *usageReporterStub) Personal(_ context.Context, request usage.PersonalRequest) (usage.PersonalReport, error) {
s.personal = request
return usage.PersonalReport{Preset: request.Preset}, s.err
}
func (s *usageReporterStub) Admin(_ context.Context, request usage.AdminRequest) (usage.AdminReport, error) {
s.admin = request
return usage.AdminReport{}, s.err
}
func usageTestHandler(t *testing.T, session identity.Session, reporter UsageReporter) http.Handler {
t.Helper()
authorizer, err := NewPlatformAuthorizer(AuthState{Required: true, Configured: true}, &fixedSessionResolver{session: session})
if err != nil {
t.Fatal(err)
}
return NewUsageHandler(authorizer, reporter, func() time.Time { return time.Date(2026, 8, 13, 0, 0, 0, 0, time.UTC) })
}