From 955f4c9182873f0896549ebf4b67c8db8d4cb6d0 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Mon, 14 Jul 2025 09:39:58 +0200
Subject: [PATCH] graph/db: refactor to make postgres fixtures re-usable

In preparation for tests where we will want to spin up SQL DBs many
times, we do some refactoring so that it is easy to re-use postgres
fixtures since those are expensive to spin up.
---
 graph/db/sql_migration_test.go |  5 +++-
 graph/db/test_postgres.go      | 45 ++++++++++++++++++++++++++++++++++
 graph/db/test_sql.go           | 23 -----------------
 graph/db/test_sqlite.go        | 33 +++++++++++++++++++++++++
 4 files changed, 82 insertions(+), 24 deletions(-)
 delete mode 100644 graph/db/test_sql.go

diff --git a/graph/db/sql_migration_test.go b/graph/db/sql_migration_test.go
index 9f32603fd..a2d8db261 100644
--- a/graph/db/sql_migration_test.go
+++ b/graph/db/sql_migration_test.go
@@ -58,6 +58,8 @@ func TestMigrateGraphToSQL(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
+	dbFixture := NewTestDBFixture(t)
+
 	writeUpdate := func(t *testing.T, db *KVStore, object any) {
 		t.Helper()
 
@@ -324,7 +326,8 @@ func TestMigrateGraphToSQL(t *testing.T) {
 			}
 
 			// Set up our destination SQL DB.
-			sql, ok := NewTestDB(t).(*SQLStore)
+			db := NewTestDBWithFixture(t, dbFixture)
+			sql, ok := db.(*SQLStore)
 			require.True(t, ok)
 
 			// Run the migration.
diff --git a/graph/db/test_postgres.go b/graph/db/test_postgres.go
index 7af420d19..5e10d94cc 100644
--- a/graph/db/test_postgres.go
+++ b/graph/db/test_postgres.go
@@ -6,9 +6,46 @@ import (
 	"database/sql"
 	"testing"
 
+	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/lightningnetwork/lnd/sqldb"
+	"github.com/stretchr/testify/require"
 )
 
+// NewTestDB is a helper function that creates a SQLStore backed by a SQL
+// database for testing.
+func NewTestDB(t testing.TB) V1Store {
+	return NewTestDBWithFixture(t, nil)
+}
+
+// NewTestDBFixture creates a new sqldb.TestPgFixture for testing purposes.
+func NewTestDBFixture(t *testing.T) *sqldb.TestPgFixture {
+	pgFixture := sqldb.NewTestPgFixture(t, sqldb.DefaultPostgresFixtureLifetime)
+	t.Cleanup(func() {
+		pgFixture.TearDown(t)
+	})
+	return pgFixture
+}
+
+// NewTestDBWithFixture is a helper function that creates a SQLStore backed by a
+// SQL database for testing.
+func NewTestDBWithFixture(t testing.TB, pgFixture *sqldb.TestPgFixture) V1Store {
+	var querier BatchedSQLQueries
+	if pgFixture == nil {
+		querier = newBatchQuerier(t)
+	} else {
+		querier = newBatchQuerierWithFixture(t, pgFixture)
+	}
+
+	store, err := NewSQLStore(
+		&SQLStoreConfig{
+			ChainHash: *chaincfg.MainNetParams.GenesisHash,
+		}, querier,
+	)
+	require.NoError(t, err)
+
+	return store
+}
+
 // newBatchQuerier creates a new BatchedSQLQueries instance for testing
 // using a PostgreSQL database fixture.
 func newBatchQuerier(t testing.TB) BatchedSQLQueries {
@@ -19,6 +56,14 @@ func newBatchQuerier(t testing.TB) BatchedSQLQueries {
 		pgFixture.TearDown(t)
 	})
 
+	return newBatchQuerierWithFixture(t, pgFixture)
+}
+
+// newBatchQuerierWithFixture creates a new BatchedSQLQueries instance for
+// testing using a PostgreSQL database fixture.
+func newBatchQuerierWithFixture(t testing.TB,
+	pgFixture *sqldb.TestPgFixture) BatchedSQLQueries {
+
 	db := sqldb.NewTestPostgresDB(t, pgFixture).BaseDB
 
 	return sqldb.NewTransactionExecutor(
diff --git a/graph/db/test_sql.go b/graph/db/test_sql.go
deleted file mode 100644
index 9d4d507b3..000000000
--- a/graph/db/test_sql.go
+++ /dev/null
@@ -1,23 +0,0 @@
-//go:build test_db_postgres || test_db_sqlite
-
-package graphdb
-
-import (
-	"testing"
-
-	"github.com/btcsuite/btcd/chaincfg"
-	"github.com/stretchr/testify/require"
-)
-
-// NewTestDB is a helper function that creates a SQLStore backed by a SQL
-// database for testing.
-func NewTestDB(t testing.TB) V1Store {
-	store, err := NewSQLStore(
-		&SQLStoreConfig{
-			ChainHash: *chaincfg.MainNetParams.GenesisHash,
-		}, newBatchQuerier(t),
-	)
-	require.NoError(t, err)
-
-	return store
-}
diff --git a/graph/db/test_sqlite.go b/graph/db/test_sqlite.go
index 35f7cb5d8..4d52b00ba 100644
--- a/graph/db/test_sqlite.go
+++ b/graph/db/test_sqlite.go
@@ -6,12 +6,45 @@ import (
 	"database/sql"
 	"testing"
 
+	"github.com/btcsuite/btcd/chaincfg"
 	"github.com/lightningnetwork/lnd/sqldb"
+	"github.com/stretchr/testify/require"
 )
 
+// NewTestDB is a helper function that creates a SQLStore backed by a SQL
+// database for testing.
+func NewTestDB(t testing.TB) V1Store {
+	return NewTestDBWithFixture(t, nil)
+}
+
+// NewTestDBFixture is a no-op for the sqlite build.
+func NewTestDBFixture(_ *testing.T) *sqldb.TestPgFixture {
+	return nil
+}
+
+// NewTestDBWithFixture is a helper function that creates a SQLStore backed by a
+// SQL database for testing.
+func NewTestDBWithFixture(t testing.TB, _ *sqldb.TestPgFixture) V1Store {
+	store, err := NewSQLStore(
+		&SQLStoreConfig{
+			ChainHash: *chaincfg.MainNetParams.GenesisHash,
+		}, newBatchQuerier(t),
+	)
+	require.NoError(t, err)
+	return store
+}
+
 // newBatchQuerier creates a new BatchedSQLQueries instance for testing
 // using a SQLite database.
 func newBatchQuerier(t testing.TB) BatchedSQLQueries {
+	return newBatchQuerierWithFixture(t, nil)
+}
+
+// newBatchQuerierWithFixture creates a new BatchedSQLQueries instance for
+// testing using a SQLite database.
+func newBatchQuerierWithFixture(t testing.TB,
+	_ *sqldb.TestPgFixture) BatchedSQLQueries {
+
 	db := sqldb.NewTestSqliteDB(t).BaseDB
 
 	return sqldb.NewTransactionExecutor(