@@ -7,38 +7,53 @@ package pprof
77
88import  (
99	"context" 
10+ 	"fmt" 
1011	"io" 
1112	"net/http" 
13+ 	"strings" 
1214	"testing" 
15+ 	"time" 
1316
1417	"github.com/stretchr/testify/require" 
18+ 
19+ 	internaltesting "github.com/envoyproxy/ai-gateway/internal/testing" 
1520)
1621
1722func  TestRun_disabled (t  * testing.T ) {
1823	t .Setenv (DisableEnvVarKey , "anything" )
1924	ctx , cancel  :=  context .WithCancel (context .Background ())
25+ 	defer  cancel ()
2026	Run (ctx )
2127	// Try accessing the pprof server here if needed. 
2228	response , err  :=  http .Get ("http://localhost:6060/debug/pprof/" ) //nolint:bodyclose 
2329	require .Error (t , err )
2430	require .Nil (t , response )
25- 	cancel ()
2631}
2732
2833func  TestRun_enabled (t  * testing.T ) {
2934	ctx , cancel  :=  context .WithCancel (context .Background ())
35+ 	defer  cancel ()
3036	Run (ctx )
31- 	// Try accessing the pprof server here if needed. 
32- 	resp , err  :=  http .Get ("http://localhost:6060/debug/pprof/cmdline" )
33- 	require .NoError (t , err )
34- 	defer  func () {
35- 		require .NoError (t , resp .Body .Close ())
36- 	}()
37- 	require .NotNil (t , resp )
38- 	body , err  :=  io .ReadAll (resp .Body )
39- 	require .NoError (t , err )
40- 	require .Contains (t , string (body ),
37+ 	// Use eventually to avoid flake when the server is not yet started by the time we access it. 
38+ 	internaltesting .RequireEventuallyNoError (t , func () error  {
39+ 		resp , err  :=  http .Get ("http://localhost:6060/debug/pprof/cmdline" )
40+ 		if  err  !=  nil  {
41+ 			return  err 
42+ 		}
43+ 		defer  func () {
44+ 			_  =  resp .Body .Close ()
45+ 		}()
46+ 		if  resp .StatusCode  !=  http .StatusOK  {
47+ 			return  fmt .Errorf ("unexpected status code: %d" , resp .StatusCode )
48+ 		}
49+ 		body , err  :=  io .ReadAll (resp .Body )
50+ 		if  err  !=  nil  {
51+ 			return  err 
52+ 		}
4153		// Test binary name should be present in the cmdline output. 
42- 		"pprof.test" )
43- 	cancel ()
54+ 		if  ! strings .Contains (string (body ), "pprof.test" ) {
55+ 			return  fmt .Errorf ("unexpected body: %s" , string (body ))
56+ 		}
57+ 		return  nil 
58+ 	}, 3 * time .Second , 100 * time .Millisecond )
4459}
0 commit comments