diff --git a/websocket_test.go b/websocket_test.go index 23d1b88cf041629fae35fec0e6dce72993fc622d..00e510c832072b21f64c163768d46eae43e3ed5e 100644 --- a/websocket_test.go +++ b/websocket_test.go @@ -485,19 +485,14 @@ func TestHandshake(t *testing.T) { } defer c.Close(websocket.StatusInternalError, "") - pctx := c.Context(ctx) + cctx := c.Context(ctx) - for ctx.Err() == nil { - err = c.Ping(ctx) - if err != nil { - if pctx.Err() == nil { - return xerrors.Errorf("context from c.Context not cancelled when connection broken") - } - return nil - } + select { + case <-ctx.Done(): + return xerrors.Errorf("child context never cancelled") + case <-cctx.Done(): + return nil } - - return xerrors.Errorf("all pings succeeded") }, }, }