Make BenchmarkDrivers work again

pkg/kv/msg_test.go

@@ -9,6 +9,7 @@ import (
 	"github.com/ory/dockertest/v3"
 	"github.com/stretchr/testify/require"
 	"github.com/treeverse/lakefs/pkg/kv"
+	_ "github.com/treeverse/lakefs/pkg/kv/dynamodb"
 	"github.com/treeverse/lakefs/pkg/kv/kvparams"
 	"github.com/treeverse/lakefs/pkg/kv/kvtest"
 	_ "github.com/treeverse/lakefs/pkg/kv/mem"

pkg/kv/store.go

@@ -167,15 +167,6 @@ func Register(name string, driver Driver) {
 	drivers[name] = driver
 }
 
-// UnregisterAllDrivers remove all loaded drivers, used for test code.
-func UnregisterAllDrivers() {
-	driversMu.Lock()
-	defer driversMu.Unlock()
-	for k := range drivers {
-		delete(drivers, k)
-	}
-}
-
 // Open lookup driver by 'type' and return store based on the configuration.
 // Failed with ErrUnknownDriver in case 'name' is not registered
 func Open(ctx context.Context, params kvparams.Config) (Store, error) {

pkg/kv/store_test.go

@@ -3,6 +3,7 @@ package kv_test
 import (
 	"context"
 	"errors"
+	"slices"
 	"sort"
 	"testing"
 
@@ -108,10 +109,13 @@ func TestRegister(t *testing.T) {
 }
 
 func TestDrivers(t *testing.T) {
-	kv.UnregisterAllDrivers()
+	preloadedDrivers := kv.Drivers()
 	kv.Register("driver1", &MockDriver{Name: "driver1"})
 	kv.Register("driver2", &MockDriver{Name: "driver2"})
 	all := kv.Drivers()
+	all = slices.DeleteFunc(all, func(name string) bool {
+		return slices.Contains(preloadedDrivers, name)
+	})
 	sort.Strings(all)
 	expectedDrivers := []string{"driver1", "driver2"}
 	if diff := deep.Equal(all, expectedDrivers); diff != nil {

pkg/testutil/dynamodb.go

@@ -93,9 +93,12 @@ func GetDynamoDBProd(ctx context.Context, tb testing.TB) kv.Store {
 	}
 	tb.Cleanup(func() {
 		defer store.Close()
+		if wrappedStore, ok := store.(*kv.StoreMetricsWrapper); ok {
+			store = wrappedStore.Store
+		}
 		err = store.(*dynamodb.Store).DropTable()
 		if err != nil {
-			tb.Fatalf("failed to delete table from DB %v %s", table, err)
+			tb.Fatalf("failed to delete table from DB %v %s", testParams.TableName, err)
 		}
 	})
 	return store