Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/server/routes/gateway_key_billing_test.go

@@ -0,0 +1,173 @@
+package routes
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/handler"
+	servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/Wei-Shaw/sub2api/internal/web"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+type keyBillingRouteAPIKeyRepo struct {
+	service.APIKeyRepository
+	apiKey *service.APIKey
+}
+
+func (r *keyBillingRouteAPIKeyRepo) GetByKeyForAuth(_ context.Context, key string) (*service.APIKey, error) {
+	if r.apiKey == nil || key != r.apiKey.Key {
+		return nil, service.ErrAPIKeyNotFound
+	}
+	clone := *r.apiKey
+	return &clone, nil
+}
+
+type keyBillingRouteRateRepo struct {
+	service.UserGroupRateRepository
+	lookupCalls int
+}
+
+func (r *keyBillingRouteRateRepo) GetByUserAndGroup(context.Context, int64, int64) (*float64, error) {
+	r.lookupCalls++
+	return nil, nil
+}
+
+func (r *keyBillingRouteRateRepo) GetRPMOverrideByUserAndGroup(context.Context, int64, int64) (*int, error) {
+	return nil, nil
+}
+
+func newKeyBillingRouteTestRouter(runMode string) (*gin.Engine, *keyBillingRouteRateRepo, string) {
+	gin.SetMode(gin.TestMode)
+	group := &service.Group{
+		ID:               42,
+		Status:           service.StatusActive,
+		Hydrated:         true,
+		Platform:         service.PlatformOpenAI,
+		SubscriptionType: service.SubscriptionTypeStandard,
+		RateMultiplier:   0.75,
+	}
+	user := &service.User{ID: 7, Role: service.RoleUser, Status: service.StatusActive, Balance: 10}
+	var groupID *int64
+	var apiKeyGroup *service.Group
+	if runMode != config.RunModeSimple {
+		groupID = &group.ID
+		apiKeyGroup = group
+	}
+	apiKey := &service.APIKey{
+		ID:      100,
+		UserID:  user.ID,
+		Key:     "billing-route-test-key",
+		Status:  service.StatusActive,
+		User:    user,
+		GroupID: groupID,
+		Group:   apiKeyGroup,
+	}
+	cfg := &config.Config{RunMode: runMode}
+	rateRepo := &keyBillingRouteRateRepo{}
+	apiKeyService := service.NewAPIKeyService(
+		&keyBillingRouteAPIKeyRepo{apiKey: apiKey}, nil, nil, nil, rateRepo, nil, cfg,
+	)
+	gatewayService := service.NewGatewayService(
+		nil, nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil, nil, nil,
+		nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
+	)
+	openAIGatewayService := service.NewOpenAIGatewayService(
+		nil, nil, nil, nil, nil, rateRepo, nil, cfg, nil, nil, nil,
+		nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
+	)
+	gatewayHandler := handler.NewGatewayHandler(
+		gatewayService, openAIGatewayService, nil, nil, nil, nil, nil, nil,
+		apiKeyService, nil, nil, nil, nil, cfg, nil,
+	)
+
+	router := gin.New()
+	if web.HasEmbeddedFrontend() {
+		router.Use(web.ServeEmbeddedFrontend())
+	}
+	RegisterGatewayRoutes(
+		router,
+		&handler.Handlers{Gateway: gatewayHandler, OpenAIGateway: &handler.OpenAIGatewayHandler{}},
+		servermiddleware.NewAPIKeyAuthMiddleware(apiKeyService, nil, cfg),
+		apiKeyService,
+		nil,
+		nil,
+		nil,
+		nil,
+		cfg,
+	)
+	return router, rateRepo, apiKey.Key
+}
+
+func TestGatewayRoutesKeyBillingInfoPathIsRegistered(t *testing.T) {
+	router := newGatewayRoutesTestRouter()
+
+	for _, route := range router.Routes() {
+		if route.Method == http.MethodGet && route.Path == "/v1/sub2api/billing" {
+			return
+		}
+	}
+
+	t.Fatal("GET /v1/sub2api/billing should be registered")
+}
+
+func TestGatewayRoutesKeyBillingInfoEndToEnd(t *testing.T) {
+	t.Run("missing credentials", func(t *testing.T) {
+		router, rateRepo, _ := newKeyBillingRouteTestRouter(config.RunModeStandard)
+		w := httptest.NewRecorder()
+
+		router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil))
+
+		require.Equal(t, http.StatusUnauthorized, w.Code)
+		require.Contains(t, w.Header().Get("Content-Type"), "application/json")
+		require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
+		require.Zero(t, rateRepo.lookupCalls)
+	})
+
+	t.Run("standard mode", func(t *testing.T) {
+		router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeStandard)
+		req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
+		req.Header.Set("Authorization", "Bearer "+key)
+		w := httptest.NewRecorder()
+
+		router.ServeHTTP(w, req)
+
+		require.Equal(t, http.StatusOK, w.Code)
+		require.Contains(t, w.Header().Get("Content-Type"), "application/json")
+		require.Equal(t, "no-store", w.Header().Get("Cache-Control"))
+		require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
+		var body map[string]any
+		require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
+		require.Equal(t, "sub2api.key_billing", body["object"])
+		require.Equal(t, 0.75, body["effective_rate_multiplier"])
+		require.Equal(t, 1, rateRepo.lookupCalls)
+	})
+
+	t.Run("simple mode", func(t *testing.T) {
+		router, rateRepo, key := newKeyBillingRouteTestRouter(config.RunModeSimple)
+		req := httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
+		req.Header.Set("x-api-key", key)
+		w := httptest.NewRecorder()
+
+		router.ServeHTTP(w, req)
+
+		require.Equal(t, http.StatusNotFound, w.Code)
+		require.Contains(t, w.Header().Get("Content-Type"), "application/json")
+		require.NotContains(t, strings.ToLower(w.Body.String()), "<!doctype html>")
+		require.JSONEq(t, `{
+			"type": "error",
+			"error": {
+				"type": "not_found_error",
+				"message": "Billing information is not supported in simple mode"
+			}
+		}`, w.Body.String())
+		require.Zero(t, rateRepo.lookupCalls)
+	})
+}