From e8c6d82fdb9d08f3b80ea0958ccfdde6af038e64 Mon Sep 17 00:00:00 2001
From: ItsT0ng <tong.niu2x@gmail.com>
Date: Sun, 1 Mar 2026 16:48:06 +1100
Subject: [PATCH] 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.
---
 pkg/utils/http_retry_test.go | 64 ++++++++++++++++++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/pkg/utils/http_retry_test.go b/pkg/utils/http_retry_test.go
index 1c2dbe115..eed654e37 100644
--- a/pkg/utils/http_retry_test.go
+++ b/pkg/utils/http_retry_test.go
@@ -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 })