Simplify usage of kadtest.CtxShort

v2/coord/coordinator_test.go

@@ -158,8 +158,7 @@ func TestConfigValidate(t *testing.T) {
 }
 
 func TestExhaustiveQuery(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -198,8 +197,7 @@ func TestExhaustiveQuery(t *testing.T) {
 }
 
 func TestRoutingUpdatedEventEmittedForCloserNodes(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -261,8 +259,7 @@ func TestRoutingUpdatedEventEmittedForCloserNodes(t *testing.T) {
 }
 
 func TestBootstrap(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -315,8 +312,7 @@ func TestBootstrap(t *testing.T) {
 }
 
 func TestIncludeNode(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)

v2/coord/network_test.go

@@ -15,8 +15,7 @@ import (
 
 // TODO: this is just a basic is-it-working test that needs to be improved
 func TestGetClosestNodes(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)

v2/coord/routing_test.go

@@ -23,8 +23,7 @@ import (
 )
 
 func TestRoutingStartBootstrapSendsEvent(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -62,8 +61,7 @@ func TestRoutingStartBootstrapSendsEvent(t *testing.T) {
 }
 
 func TestRoutingBootstrapGetClosestNodesSuccess(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -96,8 +94,7 @@ func TestRoutingBootstrapGetClosestNodesSuccess(t *testing.T) {
 }
 
 func TestRoutingBootstrapGetClosestNodesFailure(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -131,8 +128,7 @@ func TestRoutingBootstrapGetClosestNodesFailure(t *testing.T) {
 }
 
 func TestRoutingAddNodeInfoSendsEvent(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -162,8 +158,7 @@ func TestRoutingAddNodeInfoSendsEvent(t *testing.T) {
 }
 
 func TestRoutingIncludeGetClosestNodesSuccess(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -197,8 +192,7 @@ func TestRoutingIncludeGetClosestNodesSuccess(t *testing.T) {
 }
 
 func TestRoutingIncludeGetClosestNodesFailure(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)
@@ -233,8 +227,7 @@ func TestRoutingIncludeGetClosestNodesFailure(t *testing.T) {
 }
 
 func TestRoutingIncludedNodeAddToProbeList(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	clk := clock.NewMock()
 	_, nodes, err := nettest.LinearTopology(4, clk)

v2/dht_test.go

@@ -92,8 +92,7 @@ func expectEventType(t *testing.T, ctx context.Context, events <-chan coord.Rout
 }
 
 func TestAddAddresses(t *testing.T) {
-	ctx, cancel := kadtest.CtxShort(t)
-	defer cancel()
+	ctx := kadtest.CtxShort(t)
 
 	localCfg := DefaultConfig()
 

v2/internal/kadtest/context.go

@@ -6,12 +6,11 @@ import (
 	"time"
 )
 
-// CtxShort returns a Context and a CancelFunc. The context will be
-// cancelled after 10 seconds or just before the test binary deadline (as
-// specified by the -timeout flag when running the test), whichever is
-// sooner. The CancelFunc may be called to cancel the context earlier than
-// the deadline.
-func CtxShort(t *testing.T) (context.Context, context.CancelFunc) {
+// CtxShort returns a Context for tests that are expected to complete quickly.
+// The context will be cancelled after 10 seconds or just before the test
+// binary deadline (as specified by the -timeout flag when running the test), whichever
+// is sooner.
+func CtxShort(t *testing.T) context.Context {
 	t.Helper()
 
 	timeout := 10 * time.Second
@@ -27,5 +26,7 @@ func CtxShort(t *testing.T) (context.Context, context.CancelFunc) {
 		}
 	}
 
-	return context.WithDeadline(context.Background(), deadline)
+	ctx, cancel := context.WithDeadline(context.Background(), deadline)
+	t.Cleanup(cancel)
+	return ctx
 }