@@ -2,7 +2,6 @@ package executor
22
33import (
44 "context"
5- "errors"
65 "fmt"
76 "net"
87 "net/http"
@@ -241,8 +240,14 @@ func TestHTTPServerContextCancellation(t *testing.T) {
241240 errCh <- server .Start (ctx )
242241 }()
243242
244- // Give it time to start
245- time .Sleep (100 * time .Millisecond )
243+ select {
244+ case err := <- errCh :
245+ if err != nil {
246+ t .Fatalf ("Server start error: %v" , err )
247+ }
248+ case <- time .After (2 * time .Second ):
249+ t .Fatal ("Server start timed out" )
250+ }
246251
247252 // Send a request to confirm it's running
248253 client := & http.Client {Timeout : 1 * time .Second }
@@ -260,21 +265,33 @@ func TestHTTPServerContextCancellation(t *testing.T) {
260265
261266 // Cancel the context to shut down the server
262267 cancel ()
268+ client .CloseIdleConnections ()
263269
264- // Wait for shutdown to complete with timeout
265- select {
266- case err := <- errCh :
267- if err != nil && errors .Is (err , http .ErrServerClosed ) {
268- t .Fatalf ("Server shutdown error: %v" , err )
269- }
270- case <- time .After (2 * time .Second ):
271- t .Fatal ("Server shutdown timed out" )
270+ // Verify server is actually shutdown by waiting until new connections fail.
271+ shutdownClient := & http.Client {
272+ Timeout : 100 * time .Millisecond ,
273+ Transport : & http.Transport {
274+ DisableKeepAlives : true ,
275+ },
272276 }
273-
274- // Verify server is actually shutdown by attempting a new connection
275- _ , err = client .Get (fmt .Sprintf ("http://%s/store" , serverAddr ))
276- if err == nil {
277- t .Fatal ("Expected connection error after shutdown, but got none" )
277+ defer shutdownClient .CloseIdleConnections ()
278+
279+ deadline := time .After (2 * time .Second )
280+ ticker := time .NewTicker (10 * time .Millisecond )
281+ defer ticker .Stop ()
282+ for {
283+ select {
284+ case <- deadline :
285+ t .Fatal ("Expected connection error after shutdown, but server kept accepting requests" )
286+ case <- ticker .C :
287+ resp , err := shutdownClient .Get (fmt .Sprintf ("http://%s/store" , serverAddr ))
288+ if err != nil {
289+ return
290+ }
291+ if err := resp .Body .Close (); err != nil {
292+ t .Fatalf ("Failed to close response body: %v" , err )
293+ }
294+ }
278295 }
279296}
280297
0 commit comments