test(utils): add context cancellation test for DoRequestWithRetry
Verify that resp.Body is properly closed when the context is canceled during retry sleep, covering the C8 resp.Body leak fix.
This commit is contained in:
parent
53bf8bdbd0
commit
e8c6d82fdb
1 changed files with 64 additions and 0 deletions
|
|
@ -1,8 +1,11 @@
|
|||
package utils
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
|
|
@ -77,6 +80,67 @@ func TestDoRequestWithRetry(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestDoRequestWithRetry_ContextCancel(t *testing.T) {
|
||||
retryDelayUnit = 5 * time.Second // Long delay so cancel fires during sleep
|
||||
t.Cleanup(func() { retryDelayUnit = time.Second })
|
||||
|
||||
bodyClosed := false
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
w.Write([]byte("error"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
// Wrap the server's transport to detect Body.Close calls
|
||||
client := server.Client()
|
||||
client.Timeout = 30 * time.Second
|
||||
client.Transport = &bodyCloseTracker{
|
||||
rt: client.Transport,
|
||||
onClose: func() { bodyClosed = true },
|
||||
trackURL: server.URL,
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
resp, err := DoRequestWithRetry(client, req)
|
||||
require.Error(t, err, "expected error from context cancellation")
|
||||
assert.Nil(t, resp, "expected nil response when context is canceled")
|
||||
assert.True(t, bodyClosed, "expected resp.Body to be closed on context cancellation")
|
||||
}
|
||||
|
||||
// bodyCloseTracker wraps an http.RoundTripper and records when response bodies are closed.
|
||||
type bodyCloseTracker struct {
|
||||
rt http.RoundTripper
|
||||
onClose func()
|
||||
trackURL string
|
||||
}
|
||||
|
||||
func (t *bodyCloseTracker) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
resp, err := t.rt.RoundTrip(req)
|
||||
if err != nil {
|
||||
return resp, err
|
||||
}
|
||||
if strings.HasPrefix(req.URL.String(), t.trackURL) {
|
||||
resp.Body = &closeNotifier{ReadCloser: resp.Body, onClose: t.onClose}
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// closeNotifier wraps an io.ReadCloser to detect Close calls.
|
||||
type closeNotifier struct {
|
||||
io.ReadCloser
|
||||
onClose func()
|
||||
}
|
||||
|
||||
func (c *closeNotifier) Close() error {
|
||||
c.onClose()
|
||||
return c.ReadCloser.Close()
|
||||
}
|
||||
|
||||
func TestDoRequestWithRetry_Delay(t *testing.T) {
|
||||
retryDelayUnit = time.Millisecond
|
||||
t.Cleanup(func() { retryDelayUnit = time.Second })
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue