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) {
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 })
bodyClosed := false
firstRoundTripDone := make(chan struct{}, 1)
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 },
// Signal after the first round-trip response is fully constructed on the client side.
onRoundTrip: func() {
select {
case firstRoundTripDone <- struct{}{}:
default:
}
},
trackURL: server.URL,
}
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
ctx, cancel := context.WithCancel(context.Background())
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)
require.NoError(t, err)
@ -119,6 +136,7 @@ func TestDoRequestWithRetry_ContextCancel(t *testing.T) {
type bodyCloseTracker struct {
rt http.RoundTripper
onClose func()
onRoundTrip func() // called after each successful round-trip
trackURL string
}
@ -129,6 +147,9 @@ func (t *bodyCloseTracker) RoundTrip(req *http.Request) (*http.Response, error)
}
if strings.HasPrefix(req.URL.String(), t.trackURL) {
resp.Body = &closeNotifier{ReadCloser: resp.Body, onClose: t.onClose}
if t.onRoundTrip != nil {
t.onRoundTrip()
}
}
return resp, nil
}