body_test.go 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. package httputil
  2. import (
  3. "bytes"
  4. "compress/gzip"
  5. "compress/zlib"
  6. "net/http"
  7. "strings"
  8. "testing"
  9. "github.com/klauspost/compress/zstd"
  10. )
  11. const samplePayload = `{"model":"gpt-5.5","input":"hi","stream":false}`
  12. func newRequestWithBody(t *testing.T, body []byte, encoding string) *http.Request {
  13. t.Helper()
  14. req, err := http.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  15. if err != nil {
  16. t.Fatalf("NewRequest: %v", err)
  17. }
  18. if encoding != "" {
  19. req.Header.Set("Content-Encoding", encoding)
  20. }
  21. req.ContentLength = int64(len(body))
  22. return req
  23. }
  24. func TestReadRequestBodyWithPrealloc_PassesThroughIdentity(t *testing.T) {
  25. req := newRequestWithBody(t, []byte(samplePayload), "")
  26. got, err := ReadRequestBodyWithPrealloc(req)
  27. if err != nil {
  28. t.Fatalf("unexpected error: %v", err)
  29. }
  30. if string(got) != samplePayload {
  31. t.Fatalf("body mismatch: got %q", got)
  32. }
  33. }
  34. func TestReadRequestBodyWithPrealloc_DecodesZstd(t *testing.T) {
  35. enc, _ := zstd.NewWriter(nil)
  36. compressed := enc.EncodeAll([]byte(samplePayload), nil)
  37. _ = enc.Close()
  38. req := newRequestWithBody(t, compressed, "zstd")
  39. got, err := ReadRequestBodyWithPrealloc(req)
  40. if err != nil {
  41. t.Fatalf("unexpected error: %v", err)
  42. }
  43. if string(got) != samplePayload {
  44. t.Fatalf("body mismatch: got %q", got)
  45. }
  46. if req.Header.Get("Content-Encoding") != "" {
  47. t.Fatalf("Content-Encoding should be cleared after decoding")
  48. }
  49. if req.ContentLength != int64(len(samplePayload)) {
  50. t.Fatalf("ContentLength not updated: %d", req.ContentLength)
  51. }
  52. }
  53. func TestReadRequestBodyWithPrealloc_DecodesGzip(t *testing.T) {
  54. var buf bytes.Buffer
  55. gw := gzip.NewWriter(&buf)
  56. if _, err := gw.Write([]byte(samplePayload)); err != nil {
  57. t.Fatalf("gzip write: %v", err)
  58. }
  59. if err := gw.Close(); err != nil {
  60. t.Fatalf("gzip close: %v", err)
  61. }
  62. req := newRequestWithBody(t, buf.Bytes(), "gzip")
  63. got, err := ReadRequestBodyWithPrealloc(req)
  64. if err != nil {
  65. t.Fatalf("unexpected error: %v", err)
  66. }
  67. if string(got) != samplePayload {
  68. t.Fatalf("body mismatch: got %q", got)
  69. }
  70. }
  71. func TestReadRequestBodyWithPrealloc_DecodesDeflate(t *testing.T) {
  72. var buf bytes.Buffer
  73. zw := zlib.NewWriter(&buf)
  74. if _, err := zw.Write([]byte(samplePayload)); err != nil {
  75. t.Fatalf("zlib write: %v", err)
  76. }
  77. if err := zw.Close(); err != nil {
  78. t.Fatalf("zlib close: %v", err)
  79. }
  80. req := newRequestWithBody(t, buf.Bytes(), "deflate")
  81. got, err := ReadRequestBodyWithPrealloc(req)
  82. if err != nil {
  83. t.Fatalf("unexpected error: %v", err)
  84. }
  85. if string(got) != samplePayload {
  86. t.Fatalf("body mismatch: got %q", got)
  87. }
  88. }
  89. func TestReadRequestBodyWithPrealloc_RejectsUnsupportedEncoding(t *testing.T) {
  90. req := newRequestWithBody(t, []byte(samplePayload), "br")
  91. _, err := ReadRequestBodyWithPrealloc(req)
  92. if err == nil {
  93. t.Fatal("expected error for unsupported encoding, got nil")
  94. }
  95. if !strings.Contains(err.Error(), "br") {
  96. t.Fatalf("error should mention encoding, got %v", err)
  97. }
  98. }
  99. func TestReadRequestBodyWithPrealloc_RejectsCorruptZstd(t *testing.T) {
  100. req := newRequestWithBody(t, []byte("not actually zstd"), "zstd")
  101. _, err := ReadRequestBodyWithPrealloc(req)
  102. if err == nil {
  103. t.Fatal("expected error for corrupt zstd body, got nil")
  104. }
  105. }
  106. func TestReadRequestBodyWithPrealloc_NilBody(t *testing.T) {
  107. req, err := http.NewRequest(http.MethodPost, "/v1/responses", nil)
  108. if err != nil {
  109. t.Fatalf("NewRequest: %v", err)
  110. }
  111. got, err := ReadRequestBodyWithPrealloc(req)
  112. if err != nil {
  113. t.Fatalf("unexpected error: %v", err)
  114. }
  115. if got != nil {
  116. t.Fatalf("expected nil body, got %q", got)
  117. }
  118. }
  119. func TestReadRequestBodyWithPrealloc_RespectsIdentityEncoding(t *testing.T) {
  120. req := newRequestWithBody(t, []byte(samplePayload), "identity")
  121. got, err := ReadRequestBodyWithPrealloc(req)
  122. if err != nil {
  123. t.Fatalf("unexpected error: %v", err)
  124. }
  125. if string(got) != samplePayload {
  126. t.Fatalf("body mismatch: got %q", got)
  127. }
  128. }