fix(utils): eliminate flakiness in context cancellation retry test

Synchronize cancellation using an onRoundTrip callback from the
transport wrapper instead of a timing-based context timeout. This
ensures the first client.Do completes before cancel fires, so
cancellation always hits during sleepWithCtx.
This commit is contained in:
ItsT0ng 2026-03-01 16:55:34 +11:00
parent 2ab6704e59
commit e8d738cb6e

View file

@ -81,28 +81,45 @@ func TestDoRequestWithRetry(t *testing.T) {
} }
func TestDoRequestWithRetry_ContextCancel(t *testing.T) { func TestDoRequestWithRetry_ContextCancel(t *testing.T) {
retryDelayUnit = 5 * time.Second // Long delay so cancel fires during sleep // Use a long retry delay so cancellation always hits during sleepWithCtx.
retryDelayUnit = 10 * time.Second
t.Cleanup(func() { retryDelayUnit = time.Second }) t.Cleanup(func() { retryDelayUnit = time.Second })
bodyClosed := false bodyClosed := false
firstRoundTripDone := make(chan struct{}, 1)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusInternalServerError) w.WriteHeader(http.StatusInternalServerError)
w.Write([]byte("error")) w.Write([]byte("error"))
})) }))
defer server.Close() defer server.Close()
// Wrap the server's transport to detect Body.Close calls
client := server.Client() client := server.Client()
client.Timeout = 30 * time.Second client.Timeout = 30 * time.Second
client.Transport = &bodyCloseTracker{ client.Transport = &bodyCloseTracker{
rt: client.Transport, rt: client.Transport,
onClose: func() { bodyClosed = true }, onClose: func() { bodyClosed = true },
// Signal after the first round-trip response is fully constructed on the client side.
onRoundTrip: func() {
select {
case firstRoundTripDone <- struct{}{}:
default:
}
},
trackURL: server.URL, trackURL: server.URL,
} }
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
// Cancel the context after the first round-trip completes on the client side.
// This ensures client.Do has returned a valid resp (with body) and the retry
// loop is about to enter sleepWithCtx, where the cancel will be detected.
go func() {
<-firstRoundTripDone
cancel()
}()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil) req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
require.NoError(t, err) require.NoError(t, err)
@ -119,6 +136,7 @@ func TestDoRequestWithRetry_ContextCancel(t *testing.T) {
type bodyCloseTracker struct { type bodyCloseTracker struct {
rt http.RoundTripper rt http.RoundTripper
onClose func() onClose func()
onRoundTrip func() // called after each successful round-trip
trackURL string trackURL string
} }
@ -129,6 +147,9 @@ func (t *bodyCloseTracker) RoundTrip(req *http.Request) (*http.Response, error)
} }
if strings.HasPrefix(req.URL.String(), t.trackURL) { if strings.HasPrefix(req.URL.String(), t.trackURL) {
resp.Body = &closeNotifier{ReadCloser: resp.Body, onClose: t.onClose} resp.Body = &closeNotifier{ReadCloser: resp.Body, onClose: t.onClose}
if t.onRoundTrip != nil {
t.onRoundTrip()
}
} }
return resp, nil return resp, nil
} }