From 43bfcd331e33be23b4122256247cf615d4b27b64 Mon Sep 17 00:00:00 2001
From: "Leo Zhang (zhangchiqing)" <zhangchiqing@gmail.com>
Date: Fri, 10 Jan 2025 10:35:59 -0800
Subject: [PATCH] add RunWithDB

---
 storage/operation/dbtest/helper.go | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/storage/operation/dbtest/helper.go b/storage/operation/dbtest/helper.go
index af75f0da93e..058da2e17bd 100644
--- a/storage/operation/dbtest/helper.go
+++ b/storage/operation/dbtest/helper.go
@@ -21,6 +21,20 @@ func RunWithStorages(t *testing.T, fn func(*testing.T, storage.Reader, WithWrite
 	RunWithPebble(t, fn)
 }
 
+func RunWithDB(t *testing.T, fn func(*testing.T, storage.DB)) {
+	t.Run("BadgerStorage", func(t *testing.T) {
+		unittest.RunWithBadgerDB(t, func(db *badger.DB) {
+			fn(t, badgerimpl.ToDB(db))
+		})
+	})
+
+	t.Run("PebbleStorage", func(t *testing.T) {
+		unittest.RunWithPebbleDB(t, func(db *pebble.DB) {
+			fn(t, pebbleimpl.ToDB(db))
+		})
+	})
+}
+
 func RunWithBadger(t *testing.T, fn func(*testing.T, storage.Reader, WithWriter)) {
 	t.Run("BadgerStorage", func(t *testing.T) {
 		unittest.RunWithBadgerDB(t, runWithBadger(func(r storage.Reader, wr WithWriter) {