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:
parent
2ab6704e59
commit
e8d738cb6e
1 changed files with 29 additions and 8 deletions
|
|
@ -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)
|
||||||
|
|
||||||
|
|
@ -117,9 +134,10 @@ func TestDoRequestWithRetry_ContextCancel(t *testing.T) {
|
||||||
|
|
||||||
// bodyCloseTracker wraps an http.RoundTripper and records when response bodies are closed.
|
// bodyCloseTracker wraps an http.RoundTripper and records when response bodies are closed.
|
||||||
type bodyCloseTracker struct {
|
type bodyCloseTracker struct {
|
||||||
rt http.RoundTripper
|
rt http.RoundTripper
|
||||||
onClose func()
|
onClose func()
|
||||||
trackURL string
|
onRoundTrip func() // called after each successful round-trip
|
||||||
|
trackURL string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *bodyCloseTracker) RoundTrip(req *http.Request) (*http.Response, error) {
|
func (t *bodyCloseTracker) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue