Close datastore as test cleanup

Signed-off-by: sashayakovtseva <[email protected]>
sashayakovtseva · Feb 22, 2024 · f6b96c8 · f6b96c8
1 parent 36ffc9e
commit f6b96c8
Show file tree

Hide file tree

Showing 10 changed files with 49 additions and 5 deletions.
diff --git a/internal/testserver/datastore/ydb.go b/internal/testserver/datastore/ydb.go
@@ -151,11 +151,12 @@ func (y ydbTester) NewDatabase(t testing.TB) string {
 func (y ydbTester) NewDatastore(t testing.TB, initFunc InitFunc) datastore.Datastore {
 	dsn := y.NewDatabase(t)
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*15)
 	defer cancel()
 
 	migrationDriver, err := ydbMigrations.NewYDBDriver(ctx, dsn)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = migrationDriver.Close(ctx) })
 
 	err = ydbMigrations.YDBMigrations.Run(ctx, migrationDriver, migrate.Head, migrate.LiveRun)
 	require.NoError(t, err)

diff --git a/pkg/datastore/test/bulk.go b/pkg/datastore/test/bulk.go
@@ -23,6 +23,7 @@ func BulkUploadTest(t *testing.T, tester DatastoreTester) {
 
 			rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 			require.NoError(err)
+			t.Cleanup(func() { _ = rawDS.Close() })
 
 			ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
 			bulkSource := testfixtures.NewBulkTupleGenerator(
@@ -63,6 +64,7 @@ func BulkUploadErrorsTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
 

diff --git a/pkg/datastore/test/caveat.go b/pkg/datastore/test/caveat.go
@@ -29,6 +29,7 @@ func CaveatNotFoundTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx := context.Background()
 
@@ -43,6 +44,7 @@ func WriteReadDeleteCaveatTest(t *testing.T, tester DatastoreTester) {
 	req := require.New(t)
 	ds, err := tester.New(0*time.Second, veryLargeGCInterval, veryLargeGCWindow, 1)
 	req.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	skipIfNotCaveatStorer(t, ds)
 
@@ -130,6 +132,7 @@ func WriteCaveatedRelationshipTest(t *testing.T, tester DatastoreTester) {
 	req := require.New(t)
 	ds, err := tester.New(0*time.Second, veryLargeGCInterval, veryLargeGCWindow, 1)
 	req.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	skipIfNotCaveatStorer(t, ds)
 
@@ -203,6 +206,7 @@ func CaveatedRelationshipFilterTest(t *testing.T, tester DatastoreTester) {
 	req := require.New(t)
 	ds, err := tester.New(0*time.Second, veryLargeGCInterval, veryLargeGCWindow, 1)
 	req.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	skipIfNotCaveatStorer(t, ds)
 
@@ -245,6 +249,7 @@ func CaveatSnapshotReadsTest(t *testing.T, tester DatastoreTester) {
 	req := require.New(t)
 	ds, err := tester.New(0*time.Second, veryLargeGCInterval, veryLargeGCWindow, 1)
 	req.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	skipIfNotCaveatStorer(t, ds)
 
@@ -278,6 +283,7 @@ func CaveatedRelationshipWatchTest(t *testing.T, tester DatastoreTester) {
 	req := require.New(t)
 	ds, err := tester.New(0*time.Second, veryLargeGCInterval, veryLargeGCWindow, 16)
 	req.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	skipIfNotCaveatStorer(t, ds)
 	ctx, cancel := context.WithCancel(context.Background())

diff --git a/pkg/datastore/test/namespace.go b/pkg/datastore/test/namespace.go
@@ -37,6 +37,7 @@ func NamespaceNotFoundTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx := context.Background()
 
@@ -54,6 +55,7 @@ func NamespaceWriteTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx := context.Background()
 
@@ -145,6 +147,7 @@ func NamespaceDeleteTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, revision := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -193,6 +196,7 @@ func NamespaceDeleteTest(t *testing.T, tester DatastoreTester) {
 func NamespaceMultiDeleteTest(t *testing.T, tester DatastoreTester) {
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, revision := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 	ctx := context.Background()
@@ -221,6 +225,7 @@ func EmptyNamespaceDeleteTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, revision := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -266,6 +271,7 @@ definition document {
 	// Write the namespace definition to the datastore.
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx := context.Background()
 	updatedRevision, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {

diff --git a/pkg/datastore/test/pagination.go b/pkg/datastore/test/pagination.go
@@ -32,6 +32,7 @@ func OrderingTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, rev := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 	tRequire := testfixtures.TupleChecker{Require: require.New(t), DS: ds}
@@ -107,6 +108,7 @@ func LimitTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, rev := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 	tRequire := testfixtures.TupleChecker{Require: require.New(t), DS: ds}
@@ -202,6 +204,7 @@ var orderedTestCases = []struct {
 func OrderedLimitTest(t *testing.T, tester DatastoreTester) {
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, rev := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 	tRequire := testfixtures.TupleChecker{Require: require.New(t), DS: ds}
@@ -255,6 +258,7 @@ func OrderedLimitTest(t *testing.T, tester DatastoreTester) {
 func ResumeTest(t *testing.T, tester DatastoreTester) {
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, rev := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 	tRequire := testfixtures.TupleChecker{Require: require.New(t), DS: ds}
@@ -325,6 +329,7 @@ func CursorErrorsTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, rev := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
 	ctx := context.Background()

diff --git a/pkg/datastore/test/revisions.go b/pkg/datastore/test/revisions.go
@@ -34,6 +34,7 @@ func RevisionQuantizationTest(t *testing.T, tester DatastoreTester) {
 
 			ds, err := tester.New(tc.quantizationRange, veryLargeGCInterval, veryLargeGCWindow, 1)
 			require.NoError(err)
+			t.Cleanup(func() { _ = ds.Close() })
 
 			ctx := context.Background()
 			veryFirstRevision, err := ds.OptimizedRevision(ctx)
@@ -76,6 +77,7 @@ func RevisionSerializationTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
 	defer cancel()
@@ -99,6 +101,7 @@ func RevisionGCTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, 10*time.Millisecond, 300*time.Millisecond, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
@@ -178,6 +181,7 @@ func SequentialRevisionsTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, 10*time.Second, 300*time.Minute, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
@@ -201,6 +205,7 @@ func ConcurrentRevisionsTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, 10*time.Second, 300*time.Minute, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()

diff --git a/pkg/datastore/test/stats.go b/pkg/datastore/test/stats.go
@@ -18,6 +18,7 @@ func StatsTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	ds, _ = testfixtures.StandardDatastoreWithData(ds, require)
 

diff --git a/pkg/datastore/test/transactions.go b/pkg/datastore/test/transactions.go
@@ -33,6 +33,7 @@ func RetryTest(t *testing.T, tester DatastoreTester) {
 
 			rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 			require.NoError(err)
+			t.Cleanup(func() { _ = rawDS.Close() })
 
 			ds := rawDS.(TestableDatastore)
 

diff --git a/pkg/datastore/test/tuples.go b/pkg/datastore/test/tuples.go
@@ -45,7 +45,7 @@ func SimpleTest(t *testing.T, tester DatastoreTester) {
 
 			ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 			require.NoError(err)
-			defer ds.Close()
+			t.Cleanup(func() { _ = ds.Close() })
 
 			ctx := context.Background()
 
@@ -263,7 +263,7 @@ func ObjectIDsTest(t *testing.T, tester DatastoreTester) {
 
 			ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 			require.NoError(err)
-			defer ds.Close()
+			t.Cleanup(func() { _ = ds.Close() })
 
 			tpl := makeTestTuple(tc, tc)
 			require.NoError(tpl.Validate())
@@ -378,7 +378,7 @@ func DeleteRelationshipsTest(t *testing.T, tester DatastoreTester) {
 
 			ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 			require.NoError(err)
-			defer ds.Close()
+			t.Cleanup(func() { _ = ds.Close() })
 
 			setupDatastore(ds, require)
 
@@ -426,7 +426,7 @@ func InvalidReadsTest(t *testing.T, tester DatastoreTester) {
 
 		ds, err := tester.New(0, veryLargeGCInterval, testGCDuration, 1)
 		require.NoError(err)
-		defer ds.Close()
+		t.Cleanup(func() { _ = ds.Close() })
 
 		setupDatastore(ds, require)
 
@@ -470,6 +470,7 @@ func DeleteNotExistantTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -491,6 +492,7 @@ func DeleteAlreadyDeletedTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -526,6 +528,7 @@ func WriteDeleteWriteTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -553,6 +556,7 @@ func CreateAlreadyExistingTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -582,6 +586,7 @@ func TouchAlreadyExistingTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -612,6 +617,7 @@ func CreateDeleteTouchTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -641,6 +647,7 @@ func CreateTouchDeleteTouchTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -675,6 +682,7 @@ func TouchAlreadyExistingCaveatedTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -700,6 +708,7 @@ func MultipleReadsInRWTTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()
@@ -729,6 +738,7 @@ func ConcurrentWriteSerializationTest(t *testing.T, tester DatastoreTester) {
 
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = rawDS.Close() })
 
 	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require)
 	ctx := context.Background()

diff --git a/pkg/datastore/test/watch.go b/pkg/datastore/test/watch.go
@@ -60,6 +60,7 @@ func WatchTest(t *testing.T, tester DatastoreTester) {
 
 			ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 16)
 			require.NoError(err)
+			t.Cleanup(func() { _ = ds.Close() })
 
 			setupDatastore(ds, require)
 
@@ -191,6 +192,7 @@ func WatchCancelTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	startWatchRevision := setupDatastore(ds, require)
 
@@ -237,6 +239,7 @@ func WatchWithTouchTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 16)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	setupDatastore(ds, require)
 
@@ -342,6 +345,7 @@ func WatchWithDeleteTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 16)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	setupDatastore(ds, require)
 
@@ -434,6 +438,7 @@ func WatchSchemaTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 16)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	setupDatastore(ds, require)
 
@@ -494,6 +499,7 @@ func WatchAllTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 16)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	setupDatastore(ds, require)
 
@@ -626,6 +632,7 @@ func WatchCheckpointsTest(t *testing.T, tester DatastoreTester) {
 
 	ds, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 16)
 	require.NoError(err)
+	t.Cleanup(func() { _ = ds.Close() })
 
 	setupDatastore(ds, require)