diff --git a/pkg/kv/msg_test.go b/pkg/kv/msg_test.go index 98988ebc275..978ec1ae581 100644 --- a/pkg/kv/msg_test.go +++ b/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" diff --git a/pkg/kv/store.go b/pkg/kv/store.go index e5fb576db26..bdaac0bc726 100644 --- a/pkg/kv/store.go +++ b/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) { diff --git a/pkg/kv/store_test.go b/pkg/kv/store_test.go index a2c834cc4ca..afe03a708fd 100644 --- a/pkg/kv/store_test.go +++ b/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 { diff --git a/pkg/testutil/dynamodb.go b/pkg/testutil/dynamodb.go index 9768e101fac..3b85df7b5f2 100644 --- a/pkg/testutil/dynamodb.go +++ b/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