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:
ItsT0ng 2026-03-01 16:48:06 +11:00
parent 53bf8bdbd0
commit e8c6d82fdb

View file

@ -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 })