test(antigravity): add missing unit tests for upstream and custom model_mapping
- Add GetAccessToken upstream branch tests (success/failure/empty/nil) - Add mapAntigravityModel wildcard-target-equals-request edge case tests - Add upstream account smart retry test case - Add GeminiMessagesCompatService custom model_mapping and empty model tests
This commit is contained in:
@@ -265,3 +265,58 @@ func TestAntigravityGatewayService_IsModelSupported(t *testing.T) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestMapAntigravityModel_WildcardTargetEqualsRequest 测试通配符映射目标恰好等于请求模型名的 edge case
|
||||||
|
// 例如 {"claude-*": "claude-sonnet-4-5"},请求 "claude-sonnet-4-5" 时应该通过
|
||||||
|
func TestMapAntigravityModel_WildcardTargetEqualsRequest(t *testing.T) {
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
modelMapping map[string]any
|
||||||
|
requestedModel string
|
||||||
|
expected string
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "wildcard target equals request model",
|
||||||
|
modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5"},
|
||||||
|
requestedModel: "claude-sonnet-4-5",
|
||||||
|
expected: "claude-sonnet-4-5",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "wildcard target differs from request model",
|
||||||
|
modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5"},
|
||||||
|
requestedModel: "claude-opus-4-6",
|
||||||
|
expected: "claude-sonnet-4-5",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "wildcard no match",
|
||||||
|
modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5"},
|
||||||
|
requestedModel: "gpt-4o",
|
||||||
|
expected: "",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "explicit passthrough same name",
|
||||||
|
modelMapping: map[string]any{"claude-sonnet-4-5": "claude-sonnet-4-5"},
|
||||||
|
requestedModel: "claude-sonnet-4-5",
|
||||||
|
expected: "claude-sonnet-4-5",
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "multiple wildcards target equals one request",
|
||||||
|
modelMapping: map[string]any{"claude-*": "claude-sonnet-4-5", "gemini-*": "gemini-2.5-flash"},
|
||||||
|
requestedModel: "gemini-2.5-flash",
|
||||||
|
expected: "gemini-2.5-flash",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"model_mapping": tt.modelMapping,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
got := mapAntigravityModel(account, tt.requestedModel)
|
||||||
|
require.Equal(t, tt.expected, got, "mapAntigravityModel(%q) = %q, want %q", tt.requestedModel, got, tt.expected)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -526,6 +526,7 @@ func TestParseAntigravitySmartRetryInfo(t *testing.T) {
|
|||||||
func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
||||||
oauthAccount := &Account{Type: AccountTypeOAuth, Platform: PlatformAntigravity}
|
oauthAccount := &Account{Type: AccountTypeOAuth, Platform: PlatformAntigravity}
|
||||||
setupTokenAccount := &Account{Type: AccountTypeSetupToken, Platform: PlatformAntigravity}
|
setupTokenAccount := &Account{Type: AccountTypeSetupToken, Platform: PlatformAntigravity}
|
||||||
|
upstreamAccount := &Account{Type: AccountTypeUpstream, Platform: PlatformAntigravity}
|
||||||
apiKeyAccount := &Account{Type: AccountTypeAPIKey}
|
apiKeyAccount := &Account{Type: AccountTypeAPIKey}
|
||||||
|
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
@@ -587,6 +588,23 @@ func TestShouldTriggerAntigravitySmartRetry(t *testing.T) {
|
|||||||
expectedShouldRateLimit: true,
|
expectedShouldRateLimit: true,
|
||||||
modelName: "claude-sonnet-4-5",
|
modelName: "claude-sonnet-4-5",
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
name: "Upstream account with short delay - smart retry",
|
||||||
|
account: upstreamAccount,
|
||||||
|
body: `{
|
||||||
|
"error": {
|
||||||
|
"status": "RESOURCE_EXHAUSTED",
|
||||||
|
"details": [
|
||||||
|
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
|
||||||
|
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "2s"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
}`,
|
||||||
|
expectedShouldRetry: true,
|
||||||
|
expectedShouldRateLimit: false,
|
||||||
|
minWait: 2 * time.Second,
|
||||||
|
modelName: "claude-sonnet-4-5",
|
||||||
|
},
|
||||||
{
|
{
|
||||||
name: "API Key account - should not trigger",
|
name: "API Key account - should not trigger",
|
||||||
account: apiKeyAccount,
|
account: apiKeyAccount,
|
||||||
|
|||||||
97
backend/internal/service/antigravity_token_provider_test.go
Normal file
97
backend/internal/service/antigravity_token_provider_test.go
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
//go:build unit
|
||||||
|
|
||||||
|
package service
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAntigravityTokenProvider_GetAccessToken_Upstream(t *testing.T) {
|
||||||
|
provider := &AntigravityTokenProvider{}
|
||||||
|
|
||||||
|
t.Run("upstream account with valid api_key", func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Type: AccountTypeUpstream,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"api_key": "sk-test-key-12345",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), account)
|
||||||
|
require.NoError(t, err)
|
||||||
|
require.Equal(t, "sk-test-key-12345", token)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upstream account missing api_key", func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Type: AccountTypeUpstream,
|
||||||
|
Credentials: map[string]any{},
|
||||||
|
}
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), account)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "upstream account missing api_key")
|
||||||
|
require.Empty(t, token)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upstream account with empty api_key", func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Type: AccountTypeUpstream,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"api_key": "",
|
||||||
|
},
|
||||||
|
}
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), account)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "upstream account missing api_key")
|
||||||
|
require.Empty(t, token)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("upstream account with nil credentials", func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Type: AccountTypeUpstream,
|
||||||
|
}
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), account)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "upstream account missing api_key")
|
||||||
|
require.Empty(t, token)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAntigravityTokenProvider_GetAccessToken_Guards(t *testing.T) {
|
||||||
|
provider := &AntigravityTokenProvider{}
|
||||||
|
|
||||||
|
t.Run("nil account", func(t *testing.T) {
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), nil)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "account is nil")
|
||||||
|
require.Empty(t, token)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("non-antigravity platform", func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAnthropic,
|
||||||
|
Type: AccountTypeOAuth,
|
||||||
|
}
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), account)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "not an antigravity account")
|
||||||
|
require.Empty(t, token)
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("unsupported account type", func(t *testing.T) {
|
||||||
|
account := &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
}
|
||||||
|
token, err := provider.GetAccessToken(context.Background(), account)
|
||||||
|
require.Error(t, err)
|
||||||
|
require.Contains(t, err.Error(), "not an antigravity oauth account")
|
||||||
|
require.Empty(t, token)
|
||||||
|
})
|
||||||
|
}
|
||||||
@@ -905,6 +905,39 @@ func TestGeminiMessagesCompatService_isModelSupportedByAccount(t *testing.T) {
|
|||||||
model: "gpt-4",
|
model: "gpt-4",
|
||||||
expected: false,
|
expected: false,
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
name: "Antigravity平台-空模型允许",
|
||||||
|
account: &Account{Platform: PlatformAntigravity},
|
||||||
|
model: "",
|
||||||
|
expected: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Antigravity平台-自定义映射-支持自定义模型",
|
||||||
|
account: &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"model_mapping": map[string]any{
|
||||||
|
"my-custom-model": "upstream-model",
|
||||||
|
"gpt-4o": "some-model",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
model: "my-custom-model",
|
||||||
|
expected: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Antigravity平台-自定义映射-不在映射中的模型不支持",
|
||||||
|
account: &Account{
|
||||||
|
Platform: PlatformAntigravity,
|
||||||
|
Credentials: map[string]any{
|
||||||
|
"model_mapping": map[string]any{
|
||||||
|
"my-custom-model": "upstream-model",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
model: "claude-sonnet-4-5",
|
||||||
|
expected: false,
|
||||||
|
},
|
||||||
{
|
{
|
||||||
name: "Gemini平台-无映射配置-支持所有模型",
|
name: "Gemini平台-无映射配置-支持所有模型",
|
||||||
account: &Account{Platform: PlatformGemini},
|
account: &Account{Platform: PlatformGemini},
|
||||||
|
|||||||
Reference in New Issue
Block a user