李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
104 lines
3.8 KiB
Go
104 lines
3.8 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
func TestForwardResponsesInputTokensCustomRelayUsesLocalEstimate(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/input_tokens", nil)
|
|
|
|
upstream := &httpUpstreamRecorder{}
|
|
svc := &OpenAIGatewayService{
|
|
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
|
|
httpUpstream: upstream,
|
|
}
|
|
account := &Account{
|
|
ID: 159,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeAPIKey,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"api_key": "relay-key",
|
|
"base_url": "https://relay.example/v1",
|
|
},
|
|
}
|
|
body := []byte(`{"model":"gpt-5.4","instructions":"Be concise.","input":"hello world","tools":[{"type":"function","name":"lookup","description":"Look up a value","parameters":{"type":"object"}}]}`)
|
|
|
|
err := svc.ForwardResponsesInputTokens(context.Background(), c, account, body)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
require.Equal(t, "response.input_tokens", gjson.Get(recorder.Body.String(), "object").String())
|
|
require.Positive(t, gjson.Get(recorder.Body.String(), "input_tokens").Int())
|
|
require.Nil(t, upstream.lastReq, "custom relay must not receive /v1/responses/input_tokens")
|
|
}
|
|
|
|
func TestForwardResponsesInputTokensGrokOAuthUsesLocalEstimate(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/input_tokens", nil)
|
|
|
|
upstream := &httpUpstreamRecorder{}
|
|
svc := &OpenAIGatewayService{httpUpstream: upstream}
|
|
account := &Account{ID: 160, Platform: PlatformGrok, Type: AccountTypeOAuth}
|
|
body := []byte(`{"model":"grok-4.1","input":"hello world"}`)
|
|
|
|
err := svc.ForwardResponsesInputTokens(context.Background(), c, account, body)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
require.Equal(t, "response.input_tokens", gjson.Get(recorder.Body.String(), "object").String())
|
|
require.Positive(t, gjson.Get(recorder.Body.String(), "input_tokens").Int())
|
|
require.Nil(t, upstream.lastReq)
|
|
}
|
|
|
|
func TestForwardResponsesInputTokensUpstream404FallsBackLocally(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/input_tokens", nil)
|
|
|
|
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
|
StatusCode: http.StatusNotFound,
|
|
Header: make(http.Header),
|
|
Body: io.NopCloser(strings.NewReader(`{"error":{"type":"invalid_request_error","message":"Invalid URL (POST /v1/responses/input_tokens)"}}`)),
|
|
}}
|
|
svc := &OpenAIGatewayService{
|
|
cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}},
|
|
httpUpstream: upstream,
|
|
}
|
|
account := &Account{
|
|
ID: 171,
|
|
Platform: PlatformOpenAI,
|
|
Type: AccountTypeAPIKey,
|
|
Concurrency: 1,
|
|
Credentials: map[string]any{
|
|
"api_key": "official-key",
|
|
"base_url": "https://api.openai.com/v1",
|
|
},
|
|
}
|
|
body := []byte(`{"model":"gpt-5.4","instructions":"Be concise.","input":"hello world"}`)
|
|
|
|
err := svc.ForwardResponsesInputTokens(context.Background(), c, account, body)
|
|
|
|
require.NoError(t, err)
|
|
require.Equal(t, http.StatusOK, recorder.Code)
|
|
require.Equal(t, "response.input_tokens", gjson.Get(recorder.Body.String(), "object").String())
|
|
require.Positive(t, gjson.Get(recorder.Body.String(), "input_tokens").Int())
|
|
require.NotNil(t, upstream.lastReq)
|
|
}
|