plugin_manager_routing_test.go 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  1. package service
  2. import (
  3. "context"
  4. "io"
  5. "net/http"
  6. "strings"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
  9. "github.com/stretchr/testify/assert"
  10. "github.com/stretchr/testify/require"
  11. )
  12. type pluginRoutingHTTPUpstream struct {
  13. doCalls int
  14. doWithTLSCalls int
  15. }
  16. func (u *pluginRoutingHTTPUpstream) Do(*http.Request, string, int64, int) (*http.Response, error) {
  17. u.doCalls++
  18. return &http.Response{
  19. StatusCode: http.StatusOK,
  20. Header: make(http.Header),
  21. Body: io.NopCloser(strings.NewReader("legacy")),
  22. }, nil
  23. }
  24. func (u *pluginRoutingHTTPUpstream) DoWithTLS(
  25. request *http.Request,
  26. proxyURL string,
  27. accountID int64,
  28. accountConcurrency int,
  29. _ *tlsfingerprint.Profile,
  30. ) (*http.Response, error) {
  31. u.doWithTLSCalls++
  32. return u.Do(request, proxyURL, accountID, accountConcurrency)
  33. }
  34. func TestPluginManagerRoutingDoesNotTouchAPIKeyOrOtherProviders(t *testing.T) {
  35. manager := &PluginManager{}
  36. request, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.com/v1/responses", nil)
  37. require.NoError(t, err)
  38. accounts := []*Account{
  39. {ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey},
  40. {ID: 2, Platform: PlatformAnthropic, Type: AccountTypeOAuth},
  41. {ID: 3, Platform: PlatformGemini, Type: AccountTypeOAuth},
  42. }
  43. for _, account := range accounts {
  44. response, handled, routeErr := manager.RoundTripOpenAIOAuth(context.Background(), request, "", account)
  45. assert.Nil(t, response)
  46. assert.False(t, handled)
  47. assert.NoError(t, routeErr)
  48. }
  49. }
  50. func TestPluginManagerRoutingKeepsOAuthOnLegacyPathWithoutEnabledBinding(t *testing.T) {
  51. manager := &PluginManager{}
  52. request, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.com/v1/responses", nil)
  53. require.NoError(t, err)
  54. account := &Account{ID: 10, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
  55. response, handled, routeErr := manager.RoundTripOpenAIOAuth(context.Background(), request, "", account)
  56. assert.Nil(t, response)
  57. assert.False(t, handled)
  58. assert.NoError(t, routeErr)
  59. }
  60. func TestPluginManagerRoutingSelectsOnlyEligibleOpenAIOAuthAccounts(t *testing.T) {
  61. manager := &PluginManager{}
  62. manager.route.Store(&pluginRoute{pluginID: 1, rolloutPercent: 100, unavailable: "测试不可用"})
  63. assert.True(t, manager.ShouldRouteOpenAIOAuth(&Account{ID: 10, Platform: PlatformOpenAI, Type: AccountTypeOAuth}))
  64. assert.False(t, manager.ShouldRouteOpenAIOAuth(&Account{ID: 10, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}))
  65. assert.False(t, manager.ShouldRouteOpenAIOAuth(&Account{ID: 10, Platform: PlatformGrok, Type: AccountTypeOAuth}))
  66. assert.False(t, manager.ShouldRouteOpenAIOAuth(nil))
  67. }
  68. func TestOpenAIGatewayPluginRoutingPreservesAPIKeyAndFailsClosedForOAuth(t *testing.T) {
  69. manager := &PluginManager{}
  70. manager.route.Store(&pluginRoute{pluginID: 1, rolloutPercent: 100, unavailable: "测试不可用"})
  71. upstream := &pluginRoutingHTTPUpstream{}
  72. service := &OpenAIGatewayService{pluginManager: manager, httpUpstream: upstream}
  73. apiKeyRequest, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.com/v1/responses", nil)
  74. require.NoError(t, err)
  75. apiKeyResponse, err := service.doOpenAIUpstream(apiKeyRequest, "", &Account{
  76. ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Concurrency: 1,
  77. })
  78. require.NoError(t, err)
  79. require.NotNil(t, apiKeyResponse)
  80. _ = apiKeyResponse.Body.Close()
  81. assert.Equal(t, 1, upstream.doCalls)
  82. oauthRequest, err := http.NewRequestWithContext(context.Background(), http.MethodPost, "https://example.com/v1/responses", nil)
  83. require.NoError(t, err)
  84. oauthResponse, err := service.doOpenAIUpstream(oauthRequest, "", &Account{
  85. ID: 2, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1,
  86. })
  87. require.Error(t, err)
  88. assert.Nil(t, oauthResponse)
  89. assert.Contains(t, err.Error(), "插件不可用")
  90. assert.Equal(t, 1, upstream.doCalls)
  91. }
  92. func TestStablePluginBucketIsDeterministicAndBounded(t *testing.T) {
  93. for id := int64(1); id <= 1000; id++ {
  94. first := stablePluginBucket(id)
  95. assert.Equal(t, first, stablePluginBucket(id))
  96. assert.Less(t, first, uint64(100))
  97. }
  98. }