refactor(tsdb): use one test newTestDB constructor (#17638)

For tests only, we had various ways of opening DB. Reduced to one
instead of:

* Open
* newTestDB
* newTestDBOpts
* openTestDB

This so https://github.com/prometheus/prometheus/pull/17629 is smaller
and bit easier. Also for test maintainability and consistency.

Signed-off-by: bwplotka <bwplotka@gmail.com>

tsdb/compact_test.go

@@ -1257,10 +1257,7 @@ func BenchmarkCompactionFromOOOHead(b *testing.B) {
 // This is needed for unit tests that rely on
 // checking state before and after a compaction.
 func TestDisableAutoCompactions(t *testing.T) {
-	db := openTestDB(t, nil, nil)
-	defer func() {
-		require.NoError(t, db.Close())
-	}()
+	db := newTestDB(t)
 
 	blockRange := db.compactor.(*LeveledCompactor).ranges[0]
 	label := labels.FromStrings("foo", "bar")
@@ -1418,10 +1415,7 @@ func TestDeleteCompactionBlockAfterFailedReload(t *testing.T) {
 		t.Run(title, func(t *testing.T) {
 			ctx := context.Background()
 
-			db := openTestDB(t, nil, []int64{1, 100})
-			defer func() {
-				require.NoError(t, db.Close())
-			}()
+			db := newTestDB(t, withRngs(1, 100))
 			db.DisableCompactions()
 
 			expBlocks := bootStrap(db)
@@ -1993,14 +1987,11 @@ func TestDelayedCompaction(t *testing.T) {
 			}
 			t.Parallel()
 
-			var options *Options
+			var opts *Options
 			if c.compactionDelay > 0 {
-				options = &Options{CompactionDelay: c.compactionDelay}
+				opts = &Options{CompactionDelay: c.compactionDelay}
 			}
-			db := openTestDB(t, options, []int64{10})
-			defer func() {
-				require.NoError(t, db.Close())
-			}()
+			db := newTestDB(t, withOpts(opts), withRngs(10))
 
 			label := labels.FromStrings("foo", "bar")
 

tsdb/db.go

@@ -1986,6 +1986,13 @@ func (db *DB) Head() *Head {
 
 // Close the partition.
 func (db *DB) Close() error {
+	// Allow close-after-close operation for simpler use (e.g. tests).
+	select {
+	case <-db.donec:
+		return nil
+	default:
+	}
+
 	close(db.stopc)
 	if db.compactCancel != nil {
 		db.compactCancel()

tsdb/ooo_head_read_test.go

@@ -498,7 +498,7 @@ func testOOOHeadChunkReader_Chunk(t *testing.T, scenario sampleTypeScenario) {
 	minutes := func(m int64) int64 { return m * time.Minute.Milliseconds() }
 
 	t.Run("Getting a non existing chunk fails with not found error", func(t *testing.T) {
-		db := newTestDBWithOpts(t, opts)
+		db := newTestDB(t, withOpts(opts))
 
 		cr := NewHeadAndOOOChunkReader(db.head, 0, 1000, nil, nil, 0)
 		defer cr.Close()
@@ -837,7 +837,7 @@ func testOOOHeadChunkReader_Chunk(t *testing.T, scenario sampleTypeScenario) {
 
 	for _, tc := range tests {
 		t.Run(fmt.Sprintf("name=%s", tc.name), func(t *testing.T) {
-			db := newTestDBWithOpts(t, opts)
+			db := newTestDB(t, withOpts(opts))
 
 			app := db.Appender(context.Background())
 			s1Ref, _, err := scenario.appendFunc(app, s1, tc.firstInOrderSampleAt, tc.firstInOrderSampleAt/1*time.Minute.Milliseconds())
@@ -1006,7 +1006,7 @@ func testOOOHeadChunkReader_Chunk_ConsistentQueryResponseDespiteOfHeadExpanding(
 
 	for _, tc := range tests {
 		t.Run(fmt.Sprintf("name=%s", tc.name), func(t *testing.T) {
-			db := newTestDBWithOpts(t, opts)
+			db := newTestDB(t, withOpts(opts))
 
 			app := db.Appender(context.Background())
 			s1Ref, _, err := scenario.appendFunc(app, s1, tc.firstInOrderSampleAt, tc.firstInOrderSampleAt/1*time.Minute.Milliseconds())
@@ -1118,16 +1118,3 @@ func TestSortMetaByMinTimeAndMinRef(t *testing.T) {
 		})
 	}
 }
-
-func newTestDBWithOpts(t *testing.T, opts *Options) *DB {
-	dir := t.TempDir()
-
-	db, err := Open(dir, nil, nil, opts, nil)
-	require.NoError(t, err)
-
-	t.Cleanup(func() {
-		require.NoError(t, db.Close())
-	})
-
-	return db
-}

tsdb/querier_test.go

@@ -3094,11 +3094,8 @@ func TestQuerierIndexQueriesRace(t *testing.T) {
 	for _, c := range testCases {
 		t.Run(fmt.Sprintf("%v", c.matchers), func(t *testing.T) {
 			t.Parallel()
-			db := openTestDB(t, DefaultOptions(), nil)
+			db := newTestDB(t)
 			h := db.Head()
-			t.Cleanup(func() {
-				require.NoError(t, db.Close())
-			})
 			ctx, cancel := context.WithCancel(context.Background())
 			wg := &sync.WaitGroup{}
 			wg.Add(1)
@@ -3496,10 +3493,7 @@ func TestBlockBaseSeriesSet(t *testing.T) {
 }
 
 func BenchmarkHeadChunkQuerier(b *testing.B) {
-	db := openTestDB(b, nil, nil)
-	defer func() {
-		require.NoError(b, db.Close())
-	}()
+	db := newTestDB(b)
 
 	// 3h of data.
 	numTimeseries := 100
@@ -3541,10 +3535,7 @@ func BenchmarkHeadChunkQuerier(b *testing.B) {
 }
 
 func BenchmarkHeadQuerier(b *testing.B) {
-	db := openTestDB(b, nil, nil)
-	defer func() {
-		require.NoError(b, db.Close())
-	}()
+	db := newTestDB(b)
 
 	// 3h of data.
 	numTimeseries := 100
@@ -3606,12 +3597,8 @@ func TestQueryWithDeletedHistograms(t *testing.T) {
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			db := openTestDB(t, nil, nil)
-			defer func() {
-				require.NoError(t, db.Close())
-			}()
-
-			appender := db.Appender(context.Background())
+			db := newTestDB(t)
+			app := db.Appender(context.Background())
 
 			var (
 				err       error
@@ -3621,12 +3608,11 @@ func TestQueryWithDeletedHistograms(t *testing.T) {
 
 			for i := range 100 {
 				h, fh := tc(i)
-				seriesRef, err = appender.AppendHistogram(seriesRef, lbs, int64(i), h, fh)
+				seriesRef, err = app.AppendHistogram(seriesRef, lbs, int64(i), h, fh)
 				require.NoError(t, err)
 			}
 
-			err = appender.Commit()
-			require.NoError(t, err)
+			require.NoError(t, app.Commit())
 
 			matcher, err := labels.NewMatcher(labels.MatchEqual, "__name__", "test")
 			require.NoError(t, err)
@@ -3664,12 +3650,8 @@ func TestQueryWithDeletedHistograms(t *testing.T) {
 
 func TestQueryWithOneChunkCompletelyDeleted(t *testing.T) {
 	ctx := context.Background()
-	db := openTestDB(t, nil, nil)
-	defer func() {
-		require.NoError(t, db.Close())
-	}()
-
-	appender := db.Appender(context.Background())
+	db := newTestDB(t)
+	app := db.Appender(context.Background())
 
 	var (
 		err       error
@@ -3680,12 +3662,12 @@ func TestQueryWithOneChunkCompletelyDeleted(t *testing.T) {
 	// Create an int histogram chunk with samples between 0 - 20 and 30 - 40.
 	for i := range 20 {
 		h := tsdbutil.GenerateTestHistogram(1)
-		seriesRef, err = appender.AppendHistogram(seriesRef, lbs, int64(i), h, nil)
+		seriesRef, err = app.AppendHistogram(seriesRef, lbs, int64(i), h, nil)
 		require.NoError(t, err)
 	}
 	for i := 30; i < 40; i++ {
 		h := tsdbutil.GenerateTestHistogram(1)
-		seriesRef, err = appender.AppendHistogram(seriesRef, lbs, int64(i), h, nil)
+		seriesRef, err = app.AppendHistogram(seriesRef, lbs, int64(i), h, nil)
 		require.NoError(t, err)
 	}
 
@@ -3693,12 +3675,11 @@ func TestQueryWithOneChunkCompletelyDeleted(t *testing.T) {
 	// type from int histograms so a new chunk is created.
 	for i := 60; i < 100; i++ {
 		fh := tsdbutil.GenerateTestFloatHistogram(1)
-		seriesRef, err = appender.AppendHistogram(seriesRef, lbs, int64(i), nil, fh)
+		seriesRef, err = app.AppendHistogram(seriesRef, lbs, int64(i), nil, fh)
 		require.NoError(t, err)
 	}
 
-	err = appender.Commit()
-	require.NoError(t, err)
+	require.NoError(t, app.Commit())
 
 	matcher, err := labels.NewMatcher(labels.MatchEqual, "__name__", "test")
 	require.NoError(t, err)