mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-14 06:47:19 +00:00
348 lines
13 KiB
Go
348 lines
13 KiB
Go
//go:build unit
|
|
|
|
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func TestPatchGrokResponsesBodySetsMappedModelAndDropsUnsupportedFields(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
body := []byte(`{
|
|
"model": "grok",
|
|
"input": "hello",
|
|
"prompt_cache_retention": "24h",
|
|
"safety_identifier": "user-1",
|
|
"reasoning": {"effort": "high"}
|
|
}`)
|
|
|
|
patched, err := patchGrokResponsesBody(body, "grok-4.3")
|
|
require.NoError(t, err)
|
|
require.True(t, json.Valid(patched))
|
|
require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
|
|
require.False(t, gjson.GetBytes(patched, "prompt_cache_retention").Exists())
|
|
require.False(t, gjson.GetBytes(patched, "safety_identifier").Exists())
|
|
require.Equal(t, "high", gjson.GetBytes(patched, "reasoning.effort").String())
|
|
}
|
|
|
|
func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) {
|
|
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
|
|
|
account := &Account{
|
|
Platform: PlatformGrok,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"base_url": "https://xai.test/v1/",
|
|
},
|
|
}
|
|
|
|
req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token")
|
|
require.NoError(t, err)
|
|
require.Equal(t, http.MethodPost, req.Method)
|
|
require.Equal(t, "https://xai.test/v1/responses", req.URL.String())
|
|
require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
|
|
require.Equal(t, "application/json", req.Header.Get("Content-Type"))
|
|
require.Contains(t, req.Header.Get("Accept"), "text/event-stream")
|
|
|
|
data, err := io.ReadAll(req.Body)
|
|
require.NoError(t, err)
|
|
require.Equal(t, `{"model":"grok-4.3"}`, strings.TrimSpace(string(data)))
|
|
}
|
|
|
|
func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
account := &Account{
|
|
Platform: PlatformGrok,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"base_url": "https://xai.test/v1",
|
|
},
|
|
}
|
|
|
|
_, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token")
|
|
require.Error(t, err)
|
|
require.Contains(t, err.Error(), "invalid base url")
|
|
}
|
|
|
|
func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
|
|
|
account := &Account{
|
|
ID: 51,
|
|
Name: "grok",
|
|
Platform: PlatformGrok,
|
|
Type: AccountTypeOAuth,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"access_token": "access-token",
|
|
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
|
"base_url": xai.DefaultCLIBaseURL,
|
|
},
|
|
}
|
|
repo := &grokQuotaAccountRepo{
|
|
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
|
accountsByID: map[int64]*Account{51: account},
|
|
},
|
|
}
|
|
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{
|
|
"Content-Type": []string{"application/json"},
|
|
"Xai-Request-Id": []string{"xai-req"},
|
|
"X-Ratelimit-Limit-Requests": []string{"10"},
|
|
"X-Ratelimit-Remaining-Requests": []string{"9"},
|
|
"X-Ratelimit-Limit-Tokens": []string{"1000"},
|
|
"X-Ratelimit-Remaining-Tokens": []string{"990"},
|
|
},
|
|
Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}`)),
|
|
}}
|
|
svc := &OpenAIGatewayService{
|
|
httpUpstream: upstream,
|
|
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
|
accountRepo: repo,
|
|
}
|
|
|
|
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
|
require.NoError(t, err)
|
|
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
|
|
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
|
|
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
|
|
require.Equal(t, "grok", result.Model)
|
|
require.Equal(t, "grok-4.3", result.UpstreamModel)
|
|
require.Equal(t, 1, result.Usage.InputTokens)
|
|
require.Equal(t, 2, result.Usage.OutputTokens)
|
|
require.NotNil(t, repo.updates[51][grokQuotaSnapshotExtraKey])
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
}
|
|
|
|
func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
body := []byte(`{"model":"grok","input":"hi","stream":true}`)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
|
|
|
|
account := &Account{
|
|
ID: 52,
|
|
Name: "grok",
|
|
Platform: PlatformGrok,
|
|
Type: AccountTypeOAuth,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"access_token": "access-token",
|
|
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
|
"base_url": xai.DefaultCLIBaseURL,
|
|
},
|
|
}
|
|
repo := &grokQuotaAccountRepo{
|
|
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
|
accountsByID: map[int64]*Account{52: account},
|
|
},
|
|
}
|
|
upstreamBody := strings.Join([]string{
|
|
`data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`,
|
|
"",
|
|
`data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_grok","model":"grok-4.3","usage":{"input_tokens":5,"output_tokens":3,"input_tokens_details":{"cached_tokens":2}}}}`,
|
|
"",
|
|
}, "\n")
|
|
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{
|
|
"Content-Type": []string{"text/event-stream"},
|
|
"Xai-Request-Id": []string{"xai-stream-req"},
|
|
"X-Ratelimit-Limit-Requests": []string{"10"},
|
|
"X-Ratelimit-Remaining-Requests": []string{"8"},
|
|
"X-Ratelimit-Limit-Tokens": []string{"1000"},
|
|
"X-Ratelimit-Remaining-Tokens": []string{"990"},
|
|
},
|
|
Body: io.NopCloser(strings.NewReader(upstreamBody)),
|
|
}}
|
|
svc := &OpenAIGatewayService{
|
|
httpUpstream: upstream,
|
|
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
|
accountRepo: repo,
|
|
}
|
|
|
|
result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now())
|
|
require.NoError(t, err)
|
|
require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
|
|
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
|
|
require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
|
|
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
|
|
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
|
|
require.True(t, result.Stream)
|
|
require.Equal(t, "resp_grok", result.ResponseID)
|
|
require.Equal(t, "xai-stream-req", result.RequestID)
|
|
require.Equal(t, 5, result.Usage.InputTokens)
|
|
require.Equal(t, 3, result.Usage.OutputTokens)
|
|
require.Equal(t, 2, result.Usage.CacheReadInputTokens)
|
|
require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
|
|
require.Contains(t, recorder.Body.String(), "response.output_text.delta")
|
|
require.NotNil(t, repo.updates[52][grokQuotaSnapshotExtraKey])
|
|
}
|
|
|
|
func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true}`)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
account := &Account{
|
|
ID: 53,
|
|
Name: "grok",
|
|
Platform: PlatformGrok,
|
|
Type: AccountTypeOAuth,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"access_token": "access-token",
|
|
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
|
|
"base_url": xai.DefaultCLIBaseURL,
|
|
},
|
|
}
|
|
repo := &grokQuotaAccountRepo{
|
|
mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
|
|
accountsByID: map[int64]*Account{53: account},
|
|
},
|
|
}
|
|
upstreamBody := strings.Join([]string{
|
|
`data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[{"index":0,"delta":{"content":"ok"}}]}`,
|
|
"",
|
|
`data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":6,"completion_tokens":4,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":1}}}`,
|
|
"",
|
|
"data: [DONE]",
|
|
"",
|
|
}, "\n")
|
|
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
|
StatusCode: http.StatusOK,
|
|
Header: http.Header{
|
|
"Content-Type": []string{"text/event-stream"},
|
|
"X-Request-Id": []string{"chat-stream-req"},
|
|
"X-Ratelimit-Limit-Requests": []string{"10"},
|
|
"X-Ratelimit-Remaining-Requests": []string{"7"},
|
|
},
|
|
Body: io.NopCloser(strings.NewReader(upstreamBody)),
|
|
}}
|
|
svc := &OpenAIGatewayService{
|
|
cfg: rawChatCompletionsTestConfig(),
|
|
httpUpstream: upstream,
|
|
grokTokenProvider: NewGrokTokenProvider(repo, nil),
|
|
accountRepo: repo,
|
|
}
|
|
|
|
result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
|
|
require.NoError(t, err)
|
|
require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
|
|
require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
|
|
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
|
|
require.Equal(t, "sub2api-grok/1.0", upstream.lastReq.Header.Get("User-Agent"))
|
|
require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String())
|
|
require.True(t, gjson.GetBytes(upstream.lastBody, "stream_options.include_usage").Bool())
|
|
require.True(t, result.Stream)
|
|
require.Equal(t, 6, result.Usage.InputTokens)
|
|
require.Equal(t, 4, result.Usage.OutputTokens)
|
|
require.Equal(t, 1, result.Usage.CacheReadInputTokens)
|
|
require.Contains(t, recorder.Body.String(), "data: [DONE]")
|
|
require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey])
|
|
}
|
|
|
|
func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
status int
|
|
headers http.Header
|
|
wantReason string
|
|
wantMinCooldown time.Duration
|
|
wantMaxCooldown time.Duration
|
|
}{
|
|
{
|
|
name: "unauthorized reauth",
|
|
status: http.StatusUnauthorized,
|
|
wantReason: "grok oauth token unauthorized",
|
|
wantMinCooldown: 10*time.Minute - time.Second,
|
|
wantMaxCooldown: 10*time.Minute + time.Second,
|
|
},
|
|
{
|
|
name: "forbidden entitlement",
|
|
status: http.StatusForbidden,
|
|
wantReason: "grok entitlement or subscription tier denied",
|
|
wantMinCooldown: 30*time.Minute - time.Second,
|
|
wantMaxCooldown: 30*time.Minute + time.Second,
|
|
},
|
|
{
|
|
name: "rate limited retry after",
|
|
status: http.StatusTooManyRequests,
|
|
headers: http.Header{"Retry-After": []string{"45"}},
|
|
wantReason: "grok rate limited",
|
|
wantMinCooldown: 44 * time.Second,
|
|
wantMaxCooldown: 46 * time.Second,
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth}
|
|
repo := &grokQuotaAccountRepo{}
|
|
svc := &OpenAIGatewayService{accountRepo: repo}
|
|
before := time.Now()
|
|
|
|
svc.handleGrokAccountUpstreamError(context.Background(), account, tt.status, tt.headers, nil)
|
|
|
|
require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
|
|
require.Equal(t, 1, repo.tempUnschedCalls)
|
|
require.Equal(t, account.ID, repo.lastTempUnschedID)
|
|
require.Equal(t, tt.wantReason, repo.lastTempUnschedReason)
|
|
require.True(t, repo.lastTempUnschedUntil.After(before.Add(tt.wantMinCooldown)))
|
|
require.True(t, repo.lastTempUnschedUntil.Before(before.Add(tt.wantMaxCooldown)))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestHandleGrokAccountUpstreamErrorDoesNotShortenExistingPause(t *testing.T) {
|
|
existingUntil := time.Now().Add(15 * time.Minute)
|
|
account := &Account{
|
|
ID: 62,
|
|
Platform: PlatformGrok,
|
|
Type: AccountTypeOAuth,
|
|
TempUnschedulableUntil: &existingUntil,
|
|
TempUnschedulableReason: "existing pause",
|
|
}
|
|
repo := &grokQuotaAccountRepo{}
|
|
svc := &OpenAIGatewayService{accountRepo: repo}
|
|
|
|
svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil)
|
|
|
|
require.Equal(t, 1, repo.tempUnschedCalls)
|
|
require.WithinDuration(t, existingUntil, repo.lastTempUnschedUntil, time.Second)
|
|
value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
|
|
require.True(t, ok)
|
|
runtimeUntil, ok := value.(time.Time)
|
|
require.True(t, ok)
|
|
require.WithinDuration(t, existingUntil, runtimeUntil, time.Second)
|
|
}
|