Function TestMiddleware_StreamingResponse

intercept/apidump/streaming_test.go:20–96

(t *testing.T)

)

func TestMiddleware_StreamingResponse(t *testing.T) {
	t.Parallel()

	tmpDir := t.TempDir()
	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
	clk := quartz.NewMock(t)
	interceptionID := uuid.New()

	middleware := NewBridgeMiddleware(tmpDir, "openai", "gpt-4", interceptionID, logger, clk)
	require.NotNil(t, middleware)

	req, err := http.NewRequestWithContext(t.Context(), http.MethodPost, "https://api.openai.com/v1/chat/completions", bytes.NewReader([]byte(`{}`)))
	require.NoError(t, err)

	// Simulate a streaming response with multiple chunks
	chunks := []string{
		"data: {\"chunk\": 1}\n\n",
		"data: {\"chunk\": 2}\n\n",
		"data: {\"chunk\": 3}\n\n",
		"data: [DONE]\n\n",
	}

	// Create a pipe to simulate streaming
	pr, pw := io.Pipe()
	go func() {
		defer pw.Close() //nolint:revive // error handled via pipe read side
		for _, chunk := range chunks {
			if _, err := pw.Write([]byte(chunk)); err != nil {
				return
			}
		}
	}()

	resp, err := middleware(req, func(r *http.Request) (*http.Response, error) {
		return &http.Response{
			StatusCode: http.StatusOK,
			Status:     "200 OK",
			Proto:      "HTTP/1.1",
			Header:     http.Header{"Content-Type": []string{"text/event-stream"}},
			Body:       pr,
		}, nil
	})
	require.NoError(t, err)

	// Read response in small chunks to simulate streaming consumption
	var receivedData bytes.Buffer
	buf := make([]byte, 16)
	for {
		n, err := resp.Body.Read(buf)
		if n > 0 {
			_, _ = receivedData.Write(buf[:n]) // bytes.Buffer.Write never fails
		}
		if err == io.EOF {
			break
		}
		require.NoError(t, err)
	}
	require.NoError(t, resp.Body.Close())