Alanxtl commented on code in PR #706:
URL: https://github.com/apache/dubbo-go-pixiu/pull/706#discussion_r2229870783
##########
pkg/filter/llm/tokenizer/tokenizer_test.go:
##########
@@ -32,55 +35,174 @@ import (
import (
"github.com/apache/dubbo-go-pixiu/pkg/client"
+ "github.com/apache/dubbo-go-pixiu/pkg/common/constant"
"github.com/apache/dubbo-go-pixiu/pkg/context/mock"
)
-func TestUnaryResponse(t *testing.T) {
- filter := &Filter{}
-
- request, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
- assert.NoError(t, err)
- c := mock.GetMockHTTPContext(request)
- c.TargetResp = &client.UnaryResponse{
- Data: []byte(`{
- "usage": {
- "prompt_tokens": 7,
- "completion_tokens": 32,
- "total_tokens": 39,
- "prompt_tokens_details": {
- "cached_tokens": 0
+// TestUnaryResponseWithEncodings is a table-driven test for unary
(non-streaming) responses.
+// It covers multiple content encodings like gzip and deflate.
+func TestUnaryResponseWithEncodings(t *testing.T) {
+ // This is the payload we expect to process after decompression.
+ const payload = `{
+ "usage": {
+ "prompt_tokens": 7
+ }
+ }`
+
+ // Helper function to compress data with gzip for our test case.
+ compressGzipBytes := func(data string) []byte {
+ var buf bytes.Buffer
+ writer := gzip.NewWriter(&buf)
+ _, err := writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close()
+ assert.NoError(t, err)
+ return buf.Bytes()
+ }
+
+ // Helper function to compress data with flate/deflate for our test
case.
+ compressFlateBytes := func(data string) []byte {
+ var buf bytes.Buffer
+ writer, err := flate.NewWriter(&buf, -1)
+ assert.NoError(t, err)
+ _, err = writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close()
+ assert.NoError(t, err)
+ return buf.Bytes()
+ }
+
+ // Define all test cases in a table.
+ testCases := []struct {
+ name string
+ encoding string
+ getData func(string) []byte
+ }{
+ {
+ name: "No Encoding",
+ encoding: "",
+ getData: func(s string) []byte {
+ return []byte(s)
},
- "prompt_cache_hit_tokens": 0,
- "prompt_cache_miss_tokens": 7
- }
- }`)}
- filter.Encode(c)
+ },
+ {
+ name: "Gzip Encoding",
+ encoding: "gzip",
+ getData: compressGzipBytes,
+ },
+ {
+ name: "Flate Encoding",
+ encoding: "deflate",
+ getData: compressFlateBytes,
+ },
+ }
+
+ // Run the tests for each case.
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ filter := &Filter{}
+
+ request, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
+ assert.NoError(t, err)
+ c := mock.GetMockHTTPContext(request)
+
+ // Prepare the (potentially) compressed data
+ compressedData := tc.getData(payload)
+
+ c.TargetResp = &client.UnaryResponse{
+ Data: compressedData,
+ }
+ c.AddHeader(constant.HeaderKeyContentEncoding,
tc.encoding)
+
+ // Call the filter's Encode method
+ filter.Encode(c)
+ })
+ }
}
-func TestStreamResponse(t *testing.T) {
- filter := &Filter{}
-
- request, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
- assert.NoError(t, err)
- c := mock.GetMockHTTPContext(request)
- s := io.NopCloser(strings.NewReader(`data: {
- "usage": {
- "prompt_tokens": 7,
- "completion_tokens": 32,
- "total_tokens": 39,
- "prompt_tokens_details": {
- "cached_tokens": 0
+// TestStreamResponseWithEncodings is a table-driven test for streaming
responses.
+// It replaces the old TestStreamResponse.
+func TestStreamResponseWithEncodings(t *testing.T) {
+ // This is the payload we expect to process after decompression.
+ const payload = `data: {
+ "usage": {
+ "prompt_tokens": 7
+ }
+ }`
+
+ // Helper function to compress data with gzip for our test case.
+ compressGzip := func(data string) io.Reader {
+ var buf bytes.Buffer
+ writer := gzip.NewWriter(&buf)
+ _, err := writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close() // IMPORTANT: Close flushes the writer.
+ assert.NoError(t, err)
+ return &buf
+ }
+
+ compressFlate := func(data string) io.Reader {
+ var buf bytes.Buffer
+ writer, _ := flate.NewWriter(&buf, -1)
+ _, err := writer.Write([]byte(data))
+ assert.NoError(t, err)
+ err = writer.Close()
+ assert.NoError(t, err)
+ return &buf
+ }
+
+ // Define all test cases in a table.
+ testCases := []struct {
+ name string
+ encoding string
+ getStream func(string) io.Reader
+ }{
+ {
+ name: "No Encoding",
+ encoding: "",
+ getStream: func(s string) io.Reader {
+ return strings.NewReader(s)
},
- "prompt_cache_hit_tokens": 0,
- "prompt_cache_miss_tokens": 7
- }
+ },
+ {
+ name: "Gzip Encoding",
+ encoding: "gzip",
+ getStream: compressGzip,
+ },
+ {
+ name: "Flate Encoding",
+ encoding: "deflate",
+ getStream: compressFlate,
+ },
}
-`))
- c.TargetResp = &client.StreamResponse{Stream: s}
- filter.Encode(c)
- buf := make([]byte, 1024)
- c.TargetResp.(*client.StreamResponse).Stream.Read(buf)
- time.Sleep(3 * time.Millisecond)
- c.TargetResp.(*client.StreamResponse).Stream.Close()
+ // Run the tests for each case.
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ filter := &Filter{}
+
+ req, err := http.NewRequest("POST",
"http://www.dubbogopixiu.com/mock/test?name=tc";,
bytes.NewReader([]byte("{\"id\":\"12345\"}")))
+ assert.NoError(t, err)
+ ctx := mock.GetMockHTTPContext(req)
+
+ // Prepare the compressed stream and the response header
+ compressedStream := tc.getStream(payload)
+
+ // Set up the mock response
+ ctx.TargetResp = &client.StreamResponse{
+ Stream: io.NopCloser(compressedStream),
+ }
+
+ ctx.AddHeader(constant.HeaderKeyContentEncoding,
tc.encoding)
+
+ // Call the filter's Encode method
+ filter.Encode(ctx)
+
+ // Give the goroutine a moment to process the data
+ buf := make([]byte, 1024)
+ ctx.TargetResp.(*client.StreamResponse).Stream.Read(buf)
+ time.Sleep(5 * time.Millisecond)
Review Comment:
在`pkg/filter/llm/tokenizer/tokenizer.go:97`里面 对流式响应是使用goroutine处理的
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: notifications-unsubscr...@dubbo.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscr...@dubbo.apache.org
For additional commands, e-mail: notifications-h...@dubbo.apache.org
