openai_gateway_compat_cyber_test.go 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "errors"
  6. "io"
  7. "net/http"
  8. "net/http/httptest"
  9. "strings"
  10. "testing"
  11. "github.com/gin-gonic/gin"
  12. "github.com/stretchr/testify/require"
  13. )
  14. // compatCyberOAuthAccount 是 compat cyber 测试共用的 OAuth 账号。
  15. func compatCyberOAuthAccount() *Account {
  16. return &Account{
  17. ID: 1,
  18. Name: "openai-oauth",
  19. Platform: PlatformOpenAI,
  20. Type: AccountTypeOAuth,
  21. Concurrency: 1,
  22. Credentials: map[string]any{
  23. "access_token": "oauth-token",
  24. "chatgpt_account_id": "chatgpt-acc",
  25. },
  26. }
  27. }
  28. // compatCyberUpstreamSSE 构造上游 responses SSE:response.created 后 response.failed(cyber_policy)。
  29. func compatCyberUpstreamSSE() string {
  30. return strings.Join([]string{
  31. `data: {"type":"response.created","response":{"id":"resp_cyber","model":"gpt-5.5","status":"in_progress","output":[]}}`,
  32. "",
  33. `event: response.failed`,
  34. `data: {"type":"response.failed","response":{"id":"resp_cyber","object":"response","model":"gpt-5.5","status":"failed","output":[],"error":{"code":"cyber_policy","message":"flagged for cyber policy"}}}`,
  35. "",
  36. }, "\n")
  37. }
  38. func compatCyberUpstreamRecorder() *httpUpstreamRecorder {
  39. return &httpUpstreamRecorder{resp: &http.Response{
  40. StatusCode: http.StatusOK,
  41. Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid_cyber"}},
  42. Body: io.NopCloser(strings.NewReader(compatCyberUpstreamSSE())),
  43. }}
  44. }
  45. // C-1: chat completions 非流式客户端(buffered 路径)cyber 命中——不 failover、标记已设、
  46. // 以 chat 错误格式回写、丢弃 result(使 handler 落入 tokens=0 免费用量行而非 RecordUsage 扣费)。
  47. func TestForwardAsChatCompletions_BufferedCyberPolicyNoFailover(t *testing.T) {
  48. gin.SetMode(gin.TestMode)
  49. rec := httptest.NewRecorder()
  50. c, _ := gin.CreateTestContext(rec)
  51. body := []byte(`{"model":"gpt-5.5","messages":[{"role":"user","content":"hi"}],"stream":false}`)
  52. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  53. c.Request.Header.Set("Content-Type", "application/json")
  54. svc := &OpenAIGatewayService{httpUpstream: compatCyberUpstreamRecorder()}
  55. result, err := svc.ForwardAsChatCompletions(context.Background(), c, compatCyberOAuthAccount(), body, "", "gpt-5.5")
  56. require.Error(t, err)
  57. require.Nil(t, result, "cyber must drop result so handler writes tokens=0 free row, not RecordUsage")
  58. var failoverErr *UpstreamFailoverError
  59. require.False(t, errors.As(err, &failoverErr), "cyber must NOT trigger failover")
  60. mark := GetOpsCyberPolicy(c)
  61. require.NotNil(t, mark, "cyber mark must be set for handler-side recording")
  62. require.Equal(t, "cyber_policy", mark.Code)
  63. require.True(t, c.Writer.Written(), "cyber error must be written to client (passthrough)")
  64. }
  65. // I-1: chat completions 流式客户端 cyber 命中——result 必须被丢弃(返回 nil),
  66. // 使 handler forwardErrored 分支走 tokens=0 免费行,而非 RecordUsage(CyberBlocked) 扣费。
  67. func TestForwardAsChatCompletions_StreamCyberPolicyDropsResult(t *testing.T) {
  68. gin.SetMode(gin.TestMode)
  69. rec := httptest.NewRecorder()
  70. c, _ := gin.CreateTestContext(rec)
  71. body := []byte(`{"model":"gpt-5.5","messages":[{"role":"user","content":"hi"}],"stream":true}`)
  72. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  73. c.Request.Header.Set("Content-Type", "application/json")
  74. svc := &OpenAIGatewayService{httpUpstream: compatCyberUpstreamRecorder()}
  75. result, err := svc.ForwardAsChatCompletions(context.Background(), c, compatCyberOAuthAccount(), body, "", "gpt-5.5")
  76. require.Error(t, err)
  77. require.Nil(t, result, "cyber must drop result so handler does not bill via RecordUsage")
  78. var failoverErr *UpstreamFailoverError
  79. require.False(t, errors.As(err, &failoverErr), "cyber must NOT trigger failover")
  80. require.NotNil(t, GetOpsCyberPolicy(c), "cyber mark must be set")
  81. require.Contains(t, rec.Body.String(), "data: [DONE]", "stream must terminate with [DONE]")
  82. }
  83. // anthropic 非流式客户端(buffered 路径)cyber 命中——不 failover、标记已设、以 anthropic 错误格式回写、丢弃 result。
  84. func TestForwardAsAnthropic_BufferedCyberPolicyNoFailover(t *testing.T) {
  85. gin.SetMode(gin.TestMode)
  86. rec := httptest.NewRecorder()
  87. c, _ := gin.CreateTestContext(rec)
  88. body := []byte(`{"model":"gpt-5.5","max_tokens":1024,"messages":[{"role":"user","content":"hi"}],"stream":false}`)
  89. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  90. c.Request.Header.Set("Content-Type", "application/json")
  91. svc := &OpenAIGatewayService{httpUpstream: compatCyberUpstreamRecorder()}
  92. result, err := svc.ForwardAsAnthropic(context.Background(), c, compatCyberOAuthAccount(), body, "", "gpt-5.5")
  93. require.Error(t, err)
  94. require.Nil(t, result, "cyber must drop result so handler writes tokens=0 free row")
  95. var failoverErr *UpstreamFailoverError
  96. require.False(t, errors.As(err, &failoverErr), "cyber must NOT trigger failover")
  97. mark := GetOpsCyberPolicy(c)
  98. require.NotNil(t, mark, "cyber mark must be set")
  99. require.Equal(t, "cyber_policy", mark.Code)
  100. require.True(t, c.Writer.Written(), "anthropic cyber error must be written to client")
  101. require.Contains(t, rec.Body.String(), `"type":"error"`, "must use anthropic error envelope")
  102. }
  103. // anthropic 流式客户端 cyber 命中——不 failover、标记已设、下发 anthropic SSE error 事件、丢弃 result。
  104. func TestForwardAsAnthropic_StreamCyberPolicyNoFailover(t *testing.T) {
  105. gin.SetMode(gin.TestMode)
  106. rec := httptest.NewRecorder()
  107. c, _ := gin.CreateTestContext(rec)
  108. body := []byte(`{"model":"gpt-5.5","max_tokens":1024,"messages":[{"role":"user","content":"hi"}],"stream":true}`)
  109. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  110. c.Request.Header.Set("Content-Type", "application/json")
  111. svc := &OpenAIGatewayService{httpUpstream: compatCyberUpstreamRecorder()}
  112. result, err := svc.ForwardAsAnthropic(context.Background(), c, compatCyberOAuthAccount(), body, "", "gpt-5.5")
  113. require.Error(t, err)
  114. require.Nil(t, result, "cyber must drop result so handler does not bill via RecordUsage")
  115. var failoverErr *UpstreamFailoverError
  116. require.False(t, errors.As(err, &failoverErr), "cyber must NOT trigger failover")
  117. require.NotNil(t, GetOpsCyberPolicy(c), "cyber mark must be set")
  118. require.Contains(t, rec.Body.String(), "event: error", "must emit anthropic SSE error event")
  119. }