diff --git a/runtime/protocols/lib/websocket/conn_test.go b/runtime/protocols/lib/websocket/conn_test.go
index 3f0263f..acec38a 100644
--- a/runtime/protocols/lib/websocket/conn_test.go
+++ b/runtime/protocols/lib/websocket/conn_test.go
@@ -69,7 +69,8 @@
 	}
 	// Dial out in another go routine
 	go func() {
-		ctx, _ := context.RootContext()
+		ctx, cancel := context.RootContext()
+		defer cancel()
 		conn, err := WS{}.Dial(ctx, "tcp", addr.String(), time.Second)
 		numTries := 0
 		for err != nil && numTries < 5 {
diff --git a/runtime/protocols/lib/websocket/listener_test.go b/runtime/protocols/lib/websocket/listener_test.go
index 0e36a93..5aa08d1 100644
--- a/runtime/protocols/lib/websocket/listener_test.go
+++ b/runtime/protocols/lib/websocket/listener_test.go
@@ -18,7 +18,8 @@
 )
 
 func TestAcceptsAreNotSerialized(t *testing.T) {
-	ctx, _ := context.RootContext()
+	ctx, cancel := context.RootContext()
+	defer cancel()
 	ln, err := WSH{}.Listen(ctx, "wsh", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
@@ -60,7 +61,8 @@
 }
 
 func TestNonWebsocketRequest(t *testing.T) {
-	ctx, _ := context.RootContext()
+	ctx, cancel := context.RootContext()
+	defer cancel()
 	ln, err := WSH{}.Listen(ctx, "wsh", "127.0.0.1:0")
 	if err != nil {
 		t.Fatal(err)
diff --git a/runtime/protocols/lib/websocket/ws_test.go b/runtime/protocols/lib/websocket/ws_test.go
index d6325d7..cf4a1c6 100644
--- a/runtime/protocols/lib/websocket/ws_test.go
+++ b/runtime/protocols/lib/websocket/ws_test.go
@@ -48,7 +48,8 @@
 }
 
 func runTest(t *testing.T, dialObj, listenObj flow.Protocol, dialP, listenP string) {
-	ctx, _ := context.RootContext()
+	ctx, cancel := context.RootContext()
+	defer cancel()
 	address := "127.0.0.1:0"
 	timeout := time.Second
 	acceptCh := make(chan flow.Conn)