From b285546ce19a83ed26bdc5394973e993b083f2a5 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Fri, 9 May 2025 11:43:30 +0200
Subject: [PATCH] graph/db: add NewTestDB method

In this commit, we add a `test_kvdb.go` file with a single definition of
the `NewTestDB` function. A new version of `MakeTestGraph` (called
`MakeTestGraphNew` is added which makes use of this `NewTestDB` function
to create the backing `V1Store` passed to the `ChannelGraph` for tests.

Later on, we will add new implementations of this method backed by
sqlite and postgres.  When those are added, then build flags will
control which version of `NewTestDB` is called.

With this change, the only test call-site of `NewKVStore` is the new
`test_kvdb.go` file.
---
 graph/db/graph.go      | 29 +++++++++++++++++++++++++++++
 graph/db/graph_test.go | 12 ++----------
 graph/db/kv_store.go   |  3 +++
 graph/db/test_kvdb.go  | 21 +++++++++++++++++++++
 4 files changed, 55 insertions(+), 10 deletions(-)
 create mode 100644 graph/db/test_kvdb.go

diff --git a/graph/db/graph.go b/graph/db/graph.go
index 6203b2edb..cf3b3ab70 100644
--- a/graph/db/graph.go
+++ b/graph/db/graph.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sync"
 	"sync/atomic"
+	"testing"
 	"time"
 
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
@@ -13,6 +14,7 @@ import (
 	"github.com/lightningnetwork/lnd/graph/db/models"
 	"github.com/lightningnetwork/lnd/lnwire"
 	"github.com/lightningnetwork/lnd/routing/route"
+	"github.com/stretchr/testify/require"
 )
 
 // ErrChanGraphShuttingDown indicates that the ChannelGraph has shutdown or is
@@ -580,3 +582,30 @@ func (c *ChannelGraph) UpdateEdgePolicy(edge *models.ChannelEdgePolicy,
 
 	return nil
 }
+
+// MakeTestGraphNew creates a new instance of the ChannelGraph for testing
+// purposes. The backing V1Store implementation depends on the version of
+// NewTestDB included in the current build.
+//
+// NOTE: this is currently unused, but is left here for future use to show how
+// NewTestDB can be used. As the SQL implementation of the V1Store is
+// implemented, unit tests will be switched to use this function instead of
+// the existing MakeTestGraph helper. Once only this function is used, the
+// existing MakeTestGraph function will be removed and this one will be renamed.
+func MakeTestGraphNew(t testing.TB,
+	opts ...ChanGraphOption) *ChannelGraph {
+
+	t.Helper()
+
+	store := NewTestDB(t)
+
+	graph, err := NewChannelGraph(store, opts...)
+	require.NoError(t, err)
+	require.NoError(t, graph.Start())
+
+	t.Cleanup(func() {
+		require.NoError(t, graph.Stop())
+	})
+
+	return graph
+}
diff --git a/graph/db/graph_test.go b/graph/db/graph_test.go
index 9adc2ed09..62f153ede 100644
--- a/graph/db/graph_test.go
+++ b/graph/db/graph_test.go
@@ -4128,18 +4128,10 @@ func TestGraphCacheForEachNodeChannel(t *testing.T) {
 // TestGraphLoading asserts that the cache is properly reconstructed after a
 // restart.
 func TestGraphLoading(t *testing.T) {
-	// First, create a temporary directory to be used for the duration of
-	// this test.
-	tempDirName := t.TempDir()
+	t.Parallel()
 
 	// Next, create the graph for the first time.
-	backend, backendCleanup, err := kvdb.GetTestBackend(tempDirName, "cgr")
-	require.NoError(t, err)
-	defer backend.Close()
-	defer backendCleanup()
-
-	graphStore, err := NewKVStore(backend)
-	require.NoError(t, err)
+	graphStore := NewTestDB(t)
 
 	graph, err := NewChannelGraph(graphStore)
 	require.NoError(t, err)
diff --git a/graph/db/kv_store.go b/graph/db/kv_store.go
index ff221a94d..e0593cb09 100644
--- a/graph/db/kv_store.go
+++ b/graph/db/kv_store.go
@@ -4709,6 +4709,9 @@ func (c *chanGraphNodeTx) ForEachChannel(f func(*models.ChannelEdgeInfo,
 
 // MakeTestGraph creates a new instance of the ChannelGraph for testing
 // purposes.
+//
+// NOTE: this helper currently creates a ChannelGraph that is only ever backed
+// by the `KVStore` of the `V1Store` interface.
 func MakeTestGraph(t testing.TB, opts ...ChanGraphOption) *ChannelGraph {
 	t.Helper()
 
diff --git a/graph/db/test_kvdb.go b/graph/db/test_kvdb.go
new file mode 100644
index 000000000..375712972
--- /dev/null
+++ b/graph/db/test_kvdb.go
@@ -0,0 +1,21 @@
+package graphdb
+
+import (
+	"testing"
+
+	"github.com/lightningnetwork/lnd/kvdb"
+	"github.com/stretchr/testify/require"
+)
+
+// NewTestDB is a helper function that creates an BBolt database for testing.
+func NewTestDB(t testing.TB) *KVStore {
+	backend, backendCleanup, err := kvdb.GetTestBackend(t.TempDir(), "cgr")
+	require.NoError(t, err)
+
+	t.Cleanup(backendCleanup)
+
+	graphStore, err := NewKVStore(backend)
+	require.NoError(t, err)
+
+	return graphStore
+}