satellite/satellitedb/satellitedbtest: add ctx argument to Bench

Change-Id: I99a0214f4635c41ee1835ea60ad54949c78996b6
storj · Apr 23, 2024 · e9fed23 · e9fed23
1 parent 45169cb
commit e9fed23
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 15 deletions.
diff --git a/satellite/overlay/benchmark_test.go b/satellite/overlay/benchmark_test.go
@@ -17,6 +17,7 @@ import (
 	"storj.io/common/errs2"
 	"storj.io/common/pb"
 	"storj.io/common/storj"
+	"storj.io/common/testcontext"
 	"storj.io/common/testrand"
 	"storj.io/storj/satellite"
 	"storj.io/storj/satellite/nodeselection"
@@ -25,15 +26,14 @@ import (
 )
 
 func BenchmarkOverlay(b *testing.B) {
-	satellitedbtest.Bench(b, func(b *testing.B, db satellite.DB) {
+	satellitedbtest.Bench(b, func(ctx *testcontext.Context, b *testing.B, db satellite.DB) {
 		const (
 			TotalNodeCount = 211
 			OnlineCount    = 90
 			OfflineCount   = 10
 		)
 
 		overlaydb := db.OverlayCache()
-		ctx := context.Background()
 
 		var all []storj.NodeID
 		var check []storj.NodeID
@@ -179,7 +179,7 @@ func BenchmarkOverlay(b *testing.B) {
 }
 
 func BenchmarkNodeSelection(b *testing.B) {
-	satellitedbtest.Bench(b, func(b *testing.B, db satellite.DB) {
+	satellitedbtest.Bench(b, func(ctx *testcontext.Context, b *testing.B, db satellite.DB) {
 		var (
 			Total       = 10000
 			Offline     = 1000
@@ -204,8 +204,6 @@ func BenchmarkNodeSelection(b *testing.B) {
 		twoHoursAgo := now.Add(-2 * time.Hour)
 
 		overlaydb := db.OverlayCache()
-		ctx, cancel := context.WithCancel(context.Background())
-		defer cancel()
 
 		nodeSelectionConfig := overlay.NodeSelectionConfig{
 			NewNodeFraction:  newNodeFraction,

diff --git a/satellite/overlay/piececount_test.go b/satellite/overlay/piececount_test.go
@@ -86,10 +86,7 @@ func TestDB_PieceCounts(t *testing.T) {
 }
 
 func BenchmarkDB_PieceCounts(b *testing.B) {
-	satellitedbtest.Bench(b, func(b *testing.B, db satellite.DB) {
-		ctx := testcontext.New(b)
-		defer ctx.Cleanup()
-
+	satellitedbtest.Bench(b, func(ctx *testcontext.Context, b *testing.B, db satellite.DB) {
 		var NumberOfNodes = 10000
 		if testing.Short() {
 			NumberOfNodes = 1000

diff --git a/satellite/reputation/benchmark_test.go b/satellite/reputation/benchmark_test.go
@@ -4,28 +4,27 @@
 package reputation_test
 
 import (
-	"context"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/require"
 
 	"storj.io/common/storj"
+	"storj.io/common/testcontext"
 	"storj.io/common/testrand"
 	"storj.io/storj/satellite"
 	"storj.io/storj/satellite/reputation"
 	"storj.io/storj/satellite/satellitedb/satellitedbtest"
 )
 
 func BenchmarkReputation(b *testing.B) {
-	satellitedbtest.Bench(b, func(b *testing.B, db satellite.DB) {
+	satellitedbtest.Bench(b, func(ctx *testcontext.Context, b *testing.B, db satellite.DB) {
 		const (
 			TotalNodeCount = 211
 			OfflineCount   = 10
 		)
 
 		reputationdb := db.Reputation()
-		ctx := context.Background()
 
 		var all []storj.NodeID
 		for i := 0; i < TotalNodeCount; i++ {

diff --git a/satellite/satellitedb/satellitedbtest/run.go b/satellite/satellitedb/satellitedbtest/run.go
@@ -271,7 +271,7 @@ func Run(t *testing.T, test func(ctx *testcontext.Context, t *testing.T, db sate
 
 // Bench method will iterate over all supported databases. Will establish
 // connection and will create tables for each DB.
-func Bench(b *testing.B, bench func(b *testing.B, db satellite.DB)) {
+func Bench(b *testing.B, bench func(ctx *testcontext.Context, b *testing.B, db satellite.DB)) {
 	for _, dbInfo := range Databases() {
 		dbInfo := dbInfo
 		b.Run(dbInfo.Name, func(b *testing.B) {
@@ -298,8 +298,7 @@ func Bench(b *testing.B, bench func(b *testing.B, db satellite.DB)) {
 				b.Fatal(err)
 			}
 
-			// TODO: pass the ctx down
-			bench(b, db)
+			bench(ctx, b, db)
 		})
 	}
 }