mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
291 lines
9.9 KiB
Go
291 lines
9.9 KiB
Go
//go:build unit
|
|
|
|
package handler
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
|
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestGeminiV1BetaListModels_AllowlistFiltersNativeResponse(t *testing.T) {
|
|
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"next"}`)
|
|
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-2.5-pro"}})
|
|
require.True(t, ok)
|
|
require.True(t, dropped)
|
|
require.JSONEq(t, `{"models":[{"name":"models/gemini-2.5-pro"}],"nextPageToken":"next"}`, string(filtered))
|
|
}
|
|
|
|
func TestGeminiV1BetaListModels_ForcedAntigravityAppliesAllowlist(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
c.Request = httptest.NewRequest(http.MethodGet, "/antigravity/v1beta/models", nil)
|
|
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
|
Group: &service.Group{
|
|
Platform: service.PlatformGemini,
|
|
ModelAllowlist: service.GroupModelAllowlist{
|
|
Enabled: true,
|
|
Models: []string{"gemini-custom"},
|
|
},
|
|
},
|
|
})
|
|
c.Set(string(middleware.ContextKeyForcePlatform), service.PlatformAntigravity)
|
|
|
|
groupID := int64(42)
|
|
key, _ := middleware.GetAPIKeyFromContext(c)
|
|
key.GroupID = &groupID
|
|
repo := &geminiAllowlistAccountRepoStub{}
|
|
(&GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, nil, nil, nil, nil, nil)}).GeminiV1BetaListModels(c)
|
|
|
|
require.Equal(t, http.StatusOK, rec.Code)
|
|
var got antigravity.GeminiModelsListResponse
|
|
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
|
require.Empty(t, got.Models)
|
|
}
|
|
|
|
func TestGeminiModelAllowlist_DisabledPreservesNativeResponse(t *testing.T) {
|
|
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"}]}`)
|
|
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: false, Models: []string{"other"}})
|
|
require.True(t, ok)
|
|
require.False(t, dropped)
|
|
require.Equal(t, body, filtered)
|
|
}
|
|
|
|
// TestGeminiV1BetaHandler_PlatformRoutingInvariant 文档化并验证 Handler 层的平台路由逻辑不变量
|
|
// 该测试确保 gemini 和 antigravity 平台的路由逻辑符合预期
|
|
func TestGeminiV1BetaHandler_PlatformRoutingInvariant(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
platform string
|
|
expectedService string
|
|
description string
|
|
}{
|
|
{
|
|
name: "Gemini平台使用ForwardNative",
|
|
platform: service.PlatformGemini,
|
|
expectedService: "GeminiMessagesCompatService.ForwardNative",
|
|
description: "Gemini OAuth 账户直接调用 Google API",
|
|
},
|
|
{
|
|
name: "Antigravity平台使用ForwardGemini",
|
|
platform: service.PlatformAntigravity,
|
|
expectedService: "AntigravityGatewayService.ForwardGemini",
|
|
description: "Antigravity 账户通过 CRS 中转，支持 Gemini 协议",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
// 模拟 GeminiV1BetaModels 中的路由决策 (lines 199-205 in gemini_v1beta_handler.go)
|
|
var routedService string
|
|
if tt.platform == service.PlatformAntigravity {
|
|
routedService = "AntigravityGatewayService.ForwardGemini"
|
|
} else {
|
|
routedService = "GeminiMessagesCompatService.ForwardNative"
|
|
}
|
|
|
|
require.Equal(t, tt.expectedService, routedService,
|
|
"平台 %s 应该路由到 %s: %s",
|
|
tt.platform, tt.expectedService, tt.description)
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestGeminiV1BetaHandler_ListModelsAntigravityFallback 验证 ListModels 的 antigravity 降级逻辑
|
|
// 当没有 gemini 账户但有 antigravity 账户时，应返回静态模型列表
|
|
func TestGeminiV1BetaHandler_ListModelsAntigravityFallback(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
hasGeminiAccount bool
|
|
hasAntigravity bool
|
|
expectedBehavior string
|
|
}{
|
|
{
|
|
name: "有Gemini账户-调用ForwardListModels",
|
|
hasGeminiAccount: true,
|
|
hasAntigravity: false,
|
|
expectedBehavior: "forward_to_upstream",
|
|
},
|
|
{
|
|
name: "无Gemini有Antigravity-返回静态列表",
|
|
hasGeminiAccount: false,
|
|
hasAntigravity: true,
|
|
expectedBehavior: "static_fallback",
|
|
},
|
|
{
|
|
name: "无任何账户-返回503",
|
|
hasGeminiAccount: false,
|
|
hasAntigravity: false,
|
|
expectedBehavior: "service_unavailable",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
// 模拟 GeminiV1BetaListModels 的逻辑 (lines 33-44 in gemini_v1beta_handler.go)
|
|
var behavior string
|
|
|
|
if tt.hasGeminiAccount {
|
|
behavior = "forward_to_upstream"
|
|
} else if tt.hasAntigravity {
|
|
behavior = "static_fallback"
|
|
} else {
|
|
behavior = "service_unavailable"
|
|
}
|
|
|
|
require.Equal(t, tt.expectedBehavior, behavior)
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestGeminiV1BetaHandler_GetModelAntigravityFallback 验证 GetModel 的 antigravity 降级逻辑
|
|
func TestGeminiV1BetaHandler_GetModelAntigravityFallback(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
hasGeminiAccount bool
|
|
hasAntigravity bool
|
|
expectedBehavior string
|
|
}{
|
|
{
|
|
name: "有Gemini账户-调用ForwardGetModel",
|
|
hasGeminiAccount: true,
|
|
hasAntigravity: false,
|
|
expectedBehavior: "forward_to_upstream",
|
|
},
|
|
{
|
|
name: "无Gemini有Antigravity-返回静态模型信息",
|
|
hasGeminiAccount: false,
|
|
hasAntigravity: true,
|
|
expectedBehavior: "static_model_info",
|
|
},
|
|
{
|
|
name: "无任何账户-返回503",
|
|
hasGeminiAccount: false,
|
|
hasAntigravity: false,
|
|
expectedBehavior: "service_unavailable",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
// 模拟 GeminiV1BetaGetModel 的逻辑 (lines 77-87 in gemini_v1beta_handler.go)
|
|
var behavior string
|
|
|
|
if tt.hasGeminiAccount {
|
|
behavior = "forward_to_upstream"
|
|
} else if tt.hasAntigravity {
|
|
behavior = "static_model_info"
|
|
} else {
|
|
behavior = "service_unavailable"
|
|
}
|
|
|
|
require.Equal(t, tt.expectedBehavior, behavior)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestShouldFallbackGeminiModel_KnownFallbackOn404(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
res := &service.UpstreamHTTPResult{StatusCode: http.StatusNotFound}
|
|
require.True(t, shouldFallbackGeminiModel("gemini-3.1-pro-preview-customtools", res))
|
|
}
|
|
|
|
func TestShouldFallbackGeminiModel_UnknownModelOn404(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
res := &service.UpstreamHTTPResult{StatusCode: http.StatusNotFound}
|
|
require.False(t, shouldFallbackGeminiModel("gemini-future-model", res))
|
|
}
|
|
|
|
func TestShouldFallbackGeminiModel_DelegatesScopeFallback(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
res := &service.UpstreamHTTPResult{
|
|
StatusCode: http.StatusForbidden,
|
|
Headers: http.Header{"Www-Authenticate": []string{"Bearer error=\"insufficient_scope\""}},
|
|
Body: []byte("insufficient authentication scopes"),
|
|
}
|
|
require.True(t, shouldFallbackGeminiModel("gemini-future-model", res))
|
|
}
|
|
|
|
func TestGeminiV1BetaGetModel_CompositeAntigravityRoute(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
t.Run("composite route to antigravity gemini model succeeds", func(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-4-argon", nil)
|
|
req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{
|
|
Matched: true,
|
|
TargetPlatform: service.PlatformAntigravity,
|
|
UpstreamModel: "gemini-4-argon-high",
|
|
}))
|
|
c.Request = req
|
|
c.Params = gin.Params{{Key: "model", Value: "gemini-4-argon"}}
|
|
groupID := int64(4)
|
|
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
|
ID: 1,
|
|
GroupID: &groupID,
|
|
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
|
})
|
|
|
|
(&GatewayHandler{}).GeminiV1BetaGetModel(c)
|
|
require.Equal(t, http.StatusOK, rec.Code, rec.Body.String())
|
|
require.Contains(t, rec.Body.String(), "models/gemini-4-argon-high")
|
|
})
|
|
|
|
t.Run("composite route to antigravity claude model is rejected on v1beta", func(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
req := httptest.NewRequest(http.MethodGet, "/v1beta/models/claude-sonnet-5-5-high", nil)
|
|
req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{
|
|
Matched: true,
|
|
TargetPlatform: service.PlatformAntigravity,
|
|
UpstreamModel: "claude-sonnet-5-5-high",
|
|
}))
|
|
c.Request = req
|
|
c.Params = gin.Params{{Key: "model", Value: "claude-sonnet-5-5-high"}}
|
|
groupID := int64(4)
|
|
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
|
ID: 1,
|
|
GroupID: &groupID,
|
|
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
|
})
|
|
|
|
(&GatewayHandler{}).GeminiV1BetaGetModel(c)
|
|
require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
|
|
})
|
|
|
|
t.Run("composite route to antigravity claude model is rejected on generateContent", func(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(rec)
|
|
req := httptest.NewRequest(http.MethodPost, "/v1beta/models/claude-sonnet-5-5-high:generateContent", nil)
|
|
req = req.WithContext(service.WithCompositeRouteDecision(req.Context(), service.CompositeRouteDecision{
|
|
Matched: true,
|
|
TargetPlatform: service.PlatformAntigravity,
|
|
UpstreamModel: "claude-sonnet-5-5-high",
|
|
}))
|
|
c.Request = req
|
|
c.Params = gin.Params{{Key: "modelAction", Value: "/claude-sonnet-5-5-high:generateContent"}}
|
|
groupID := int64(4)
|
|
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
|
ID: 1,
|
|
GroupID: &groupID,
|
|
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
|
})
|
|
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1, Concurrency: 1})
|
|
|
|
(&GatewayHandler{}).GeminiV1BetaModels(c)
|
|
require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String())
|
|
})
|
|
}
|