openai_apikey_responses_probe_test.go 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161
  1. package service
  2. import (
  3. "context"
  4. "io"
  5. "net/http"
  6. "strings"
  7. "testing"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat"
  11. "github.com/stretchr/testify/require"
  12. )
  13. func TestProbeOpenAIAPIKeyResponsesSupportUsesCodexProbeHeaders(t *testing.T) {
  14. updateCalls := make(chan map[string]any, 1)
  15. account := Account{
  16. ID: 96,
  17. Platform: PlatformOpenAI,
  18. Type: AccountTypeAPIKey,
  19. Concurrency: 1,
  20. Credentials: map[string]any{
  21. "api_key": "sk-test",
  22. "base_url": "https://compat-upstream.example/v1",
  23. },
  24. }
  25. repo := &snapshotUpdateAccountRepo{
  26. stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
  27. updateExtraCalls: updateCalls,
  28. }
  29. upstream := &httpUpstreamRecorder{resp: &http.Response{
  30. StatusCode: http.StatusOK,
  31. Header: make(http.Header),
  32. Body: io.NopCloser(strings.NewReader(`{"output":[{"type":"function_call","name":"probe_ping"}]}`)),
  33. }}
  34. svc := &AccountTestService{
  35. accountRepo: repo,
  36. httpUpstream: upstream,
  37. cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
  38. }
  39. svc.ProbeOpenAIAPIKeyResponsesSupport(context.Background(), account.ID)
  40. require.NotNil(t, upstream.lastReq)
  41. require.Equal(t, "https://compat-upstream.example/v1/responses", upstream.lastReq.URL.String())
  42. requireOpenAICodexProbeHeaders(t, upstream.lastReq.Header)
  43. updates := <-updateCalls
  44. require.Equal(t, true, updates[openai_compat.ExtraKeyResponsesSupported])
  45. }
  46. func TestProbeOpenAIAPIKeyResponsesSupportCNProviders(t *testing.T) {
  47. tests := []struct {
  48. name string
  49. id int64
  50. platform string
  51. protocol string
  52. wantSupport bool
  53. wantMode string
  54. }{
  55. {name: "deepseek adaptive supports responses", id: 201, platform: PlatformDeepseek, protocol: APIProtocolAdaptive, wantSupport: true, wantMode: string(openai_compat.ResponsesSupportModeForceResponses)},
  56. {name: "deepseek chat clears forced responses", id: 202, platform: PlatformDeepseek, protocol: APIProtocolChatCompletions, wantSupport: false, wantMode: string(openai_compat.ResponsesSupportModeAuto)},
  57. {name: "kimi adaptive supports responses", id: 203, platform: PlatformKimi, protocol: APIProtocolAdaptive, wantSupport: true, wantMode: string(openai_compat.ResponsesSupportModeForceResponses)},
  58. {name: "kimi responses protocol supports responses", id: 205, platform: PlatformKimi, protocol: APIProtocolResponses, wantSupport: true, wantMode: string(openai_compat.ResponsesSupportModeForceResponses)},
  59. {name: "zhipu adaptive falls back to chat", id: 204, platform: PlatformZhipu, protocol: APIProtocolAdaptive, wantSupport: false, wantMode: string(openai_compat.ResponsesSupportModeAuto)},
  60. }
  61. for _, tc := range tests {
  62. t.Run(tc.name, func(t *testing.T) {
  63. updateCalls := make(chan map[string]any, 1)
  64. account := Account{
  65. ID: tc.id, Platform: tc.platform, Type: AccountTypeAPIKey,
  66. Credentials: map[string]any{"api_key": "sk-test", "api_protocol": tc.protocol},
  67. Extra: map[string]any{
  68. openai_compat.ExtraKeyResponsesMode: string(openai_compat.ResponsesSupportModeForceResponses),
  69. },
  70. }
  71. repo := &snapshotUpdateAccountRepo{
  72. stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
  73. updateExtraCalls: updateCalls,
  74. }
  75. svc := &AccountTestService{accountRepo: repo}
  76. svc.ProbeOpenAIAPIKeyResponsesSupport(context.Background(), account.ID)
  77. updates := <-updateCalls
  78. require.Equal(t, tc.wantSupport, updates[openai_compat.ExtraKeyResponsesSupported])
  79. require.Equal(t, tc.wantMode, updates[openai_compat.ExtraKeyResponsesMode])
  80. })
  81. }
  82. }
  83. func TestDecideResponsesProbeSupport(t *testing.T) {
  84. fnCall := []byte(`{"output":[{"type":"reasoning"},{"type":"function_call","name":"probe_ping"}]}`)
  85. reasoningOnly := []byte(`{"output":[{"type":"reasoning"}]}`)
  86. cases := []struct {
  87. name string
  88. status int
  89. body []byte
  90. want bool
  91. }{
  92. // Endpoint clearly absent on third-party OpenAI-compatible upstreams.
  93. {"404 endpoint absent", 404, fnCall, false},
  94. {"405 method not allowed", 405, fnCall, false},
  95. // 2xx: tool capability is judged by presence of a function_call output item.
  96. {"200 with function_call", 200, fnCall, true},
  97. // Volcengine Ark coding/v3 × kimi-k2.6: reasoning only, no function_call.
  98. {"200 reasoning only", 200, reasoningOnly, false},
  99. {"200 invalid json", 200, []byte("not-json"), false},
  100. {"200 no output field", 200, []byte(`{"status":"completed"}`), false},
  101. // Non-2xx (other than 404/405): endpoint exists, capability undecidable -> conservative true.
  102. {"400 conservative true", 400, reasoningOnly, true},
  103. {"401 conservative true", 401, nil, true},
  104. {"500 conservative true", 500, nil, true},
  105. }
  106. for _, tc := range cases {
  107. t.Run(tc.name, func(t *testing.T) {
  108. require.Equal(t, tc.want, decideResponsesProbeSupport(tc.status, tc.body))
  109. })
  110. }
  111. }
  112. func TestResponsesProbeBodyHasFunctionCall(t *testing.T) {
  113. require.True(t, responsesProbeBodyHasFunctionCall([]byte(`{"output":[{"type":"function_call"}]}`)))
  114. require.True(t, responsesProbeBodyHasFunctionCall([]byte(`{"output":[{"type":"reasoning"},{"type":"function_call"}]}`)))
  115. require.False(t, responsesProbeBodyHasFunctionCall([]byte(`{"output":[{"type":"reasoning"}]}`)))
  116. require.False(t, responsesProbeBodyHasFunctionCall([]byte(`{"output":[]}`)))
  117. require.False(t, responsesProbeBodyHasFunctionCall([]byte(`{}`)))
  118. require.False(t, responsesProbeBodyHasFunctionCall([]byte(`garbage`)))
  119. }
  120. func TestSelectResponsesProbeModel(t *testing.T) {
  121. // No model_mapping -> fall back to DefaultTestModel (OpenAI official APIKey).
  122. require.Equal(t, openai.DefaultTestModel, selectResponsesProbeModel(&Account{}))
  123. // model_mapping values are upstream models; pick first by sort for reproducibility.
  124. acct := &Account{Credentials: map[string]any{
  125. "model_mapping": map[string]any{
  126. "client-b": "zeta-model",
  127. "client-a": "alpha-model",
  128. },
  129. }}
  130. require.Equal(t, "alpha-model", selectResponsesProbeModel(acct))
  131. // Wildcard / blank upstream values are skipped.
  132. acctWild := &Account{Credentials: map[string]any{
  133. "model_mapping": map[string]any{
  134. "a": "*",
  135. "b": " ",
  136. "c": "real-model",
  137. },
  138. }}
  139. require.Equal(t, "real-model", selectResponsesProbeModel(acctWild))
  140. // Only wildcard mappings -> DefaultTestModel.
  141. acctAllWild := &Account{Credentials: map[string]any{
  142. "model_mapping": map[string]any{"a": "gpt-*"},
  143. }}
  144. require.Equal(t, openai.DefaultTestModel, selectResponsesProbeModel(acctAllWild))
  145. }