openai_response_terminal_usage_compat_test.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. //go:build unit
  2. package service
  3. import (
  4. "strings"
  5. "testing"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  8. "github.com/stretchr/testify/require"
  9. "github.com/tidwall/gjson"
  10. )
  11. func TestEffectiveOpenAISSEEventTypePrefersPayload(t *testing.T) {
  12. t.Parallel()
  13. require.Equal(t, "response.failed", effectiveOpenAISSEEventType([]byte(`{"type":"response.failed"}`), "error"))
  14. require.Equal(t, "error", effectiveOpenAISSEEventType([]byte(`{"error":{"message":"failed"}}`), " error "))
  15. require.JSONEq(t, `{"type":"error","error":{"message":"failed"}}`, openAICompatPayloadWithEventType(`{"type":"","error":{"message":"failed"}}`, "error"))
  16. }
  17. func TestExtractOpenAISSETerminalEventUsesEventField(t *testing.T) {
  18. t.Parallel()
  19. body := "event: error\n" +
  20. "data: {\"error\":{\"message\":\"provider failed\"}}\n\n" +
  21. "data: [DONE]\n\n"
  22. eventType, payload, ok := extractOpenAISSETerminalEvent(body)
  23. require.True(t, ok)
  24. require.Equal(t, "error", eventType)
  25. require.Equal(t, "provider failed", extractOpenAISSEErrorMessage(payload))
  26. }
  27. func TestExtractOpenAISSETerminalEventUsesFinalAuthoritativeTerminal(t *testing.T) {
  28. t.Parallel()
  29. body := "data: {\"type\":\"error\",\"error\":{\"message\":\"recovering\"}}\n\n" +
  30. "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_1\",\"status\":\"completed\"}}\n\n"
  31. eventType, payload, ok := extractOpenAISSETerminalEvent(body)
  32. require.True(t, ok)
  33. require.Equal(t, "response.completed", eventType)
  34. require.Equal(t, "resp_1", gjson.GetBytes(payload, "response.id").String())
  35. }
  36. func TestParseSSEUsageEffectiveTerminalRules(t *testing.T) {
  37. t.Parallel()
  38. svc := &OpenAIGatewayService{}
  39. usage := &OpenAIUsage{}
  40. svc.parseSSEUsageBytesWithType([]byte(`{"usage":{"input_tokens":17,"output_tokens":5,"input_tokens_details":{"cached_tokens":3}}}`), "response.in_progress", usage)
  41. svc.parseSSEUsageBytesWithType([]byte(`{"response":{"id":"resp_1"}}`), "response.completed", usage)
  42. require.Equal(t, OpenAIUsage{InputTokens: 17, OutputTokens: 5, CacheReadInputTokens: 3}, *usage)
  43. svc.parseSSEUsageBytesWithType([]byte(`{"response":{"usage":{"input_tokens":0,"output_tokens":0,"input_tokens_details":{"cached_tokens":0}}}}`), "response.completed", usage)
  44. require.Equal(t, OpenAIUsage{InputTokens: 17, OutputTokens: 5, CacheReadInputTokens: 3}, *usage)
  45. svc.parseSSEUsageBytesWithType([]byte(`{"response":{"usage":{"input_tokens":2,"output_tokens":0,"input_tokens_details":{"cached_tokens":0}}}}`), "response.completed", usage)
  46. require.Equal(t, OpenAIUsage{InputTokens: 2}, *usage)
  47. }
  48. func BenchmarkParseSSEUsageNoUsageDelta(b *testing.B) {
  49. svc := &OpenAIGatewayService{}
  50. usage := &OpenAIUsage{}
  51. payload := []byte(`{"type":"response.output_text.delta","delta":"hello"}`)
  52. b.ReportAllocs()
  53. b.ResetTimer()
  54. for i := 0; i < b.N; i++ {
  55. svc.parseSSEUsageBytesWithType(payload, "response.output_text.delta", usage)
  56. }
  57. }
  58. func TestOpenAICompatTerminalResponseSynthesizesBareError(t *testing.T) {
  59. t.Parallel()
  60. payload := []byte(`{"type":"error","code":"upstream_error","error":{"message":"provider failed"}}`)
  61. event := &apicompat.ResponsesStreamEvent{Type: "error", Code: "upstream_error"}
  62. response := openAICompatTerminalResponse(event, payload)
  63. require.NotNil(t, response)
  64. require.Equal(t, "failed", response.Status)
  65. require.Equal(t, "upstream_error", response.Error.Code)
  66. require.Equal(t, "provider failed", response.Error.Message)
  67. }
  68. func TestForEachOpenAISSEFrameDataTypeOverridesEventField(t *testing.T) {
  69. t.Parallel()
  70. var types []string
  71. forEachOpenAISSEFrame(strings.Join([]string{
  72. "event: response.in_progress",
  73. `data: {"type":"response.completed","response":{"id":"resp_1"}}`,
  74. "",
  75. }, "\n"), func(eventType string, _ []byte) {
  76. types = append(types, eventType)
  77. })
  78. require.Equal(t, []string{"response.completed"}, types)
  79. }
  80. func TestSampleOpenAIMissingUsageLogRateLimitsAndCounts(t *testing.T) {
  81. var sampler openAIMissingUsageLogSampler
  82. now := time.Unix(100, 0)
  83. logNow, total, suppressed := sampler.sample(now)
  84. require.True(t, logNow)
  85. require.Equal(t, uint64(1), total)
  86. require.Zero(t, suppressed)
  87. logNow, total, _ = sampler.sample(now.Add(time.Second))
  88. require.False(t, logNow)
  89. require.Equal(t, uint64(2), total)
  90. logNow, total, suppressed = sampler.sample(now.Add(openAIMissingUsageLogInterval))
  91. require.True(t, logNow)
  92. require.Equal(t, uint64(3), total)
  93. require.Equal(t, uint64(1), suppressed)
  94. }