diff --git a/channeldb/db_test.go b/channeldb/db_test.go
index e175ea1fb..9c8c34d58 100644
--- a/channeldb/db_test.go
+++ b/channeldb/db_test.go
@@ -14,7 +14,6 @@ import (
 	"github.com/btcsuite/btcd/btcutil"
 	"github.com/btcsuite/btcd/chaincfg/chainhash"
 	"github.com/btcsuite/btcd/wire"
-	graphdb "github.com/lightningnetwork/lnd/graph/db"
 	"github.com/lightningnetwork/lnd/graph/db/models"
 	"github.com/lightningnetwork/lnd/keychain"
 	"github.com/lightningnetwork/lnd/kvdb"
@@ -185,21 +184,22 @@ func TestMultiSourceAddrsForNode(t *testing.T) {
 	fullDB, err := MakeTestDB(t)
 	require.NoError(t, err, "unable to make test database")
 
-	graph, err := graphdb.MakeTestGraph(t)
-	require.NoError(t, err)
+	graph := newMockAddrSource(t)
+	t.Cleanup(func() {
+		graph.AssertExpectations(t)
+	})
 
-	// We'll make a test vertex to insert into the database, as the source
-	// node, but this node will only have half the number of addresses it
-	// usually does.
+	// We'll make a test vertex, but this node will only have half the
+	// number of addresses it usually does.
 	testNode := createTestVertex(t)
-	require.NoError(t, err, "unable to create test node")
-	testNode.Addresses = []net.Addr{testAddr}
-	require.NoError(t, graph.SetSourceNode(testNode))
+	nodePub, err := testNode.PubKey()
+	require.NoError(t, err)
+	graph.On("AddrsForNode", nodePub).Return(
+		true, []net.Addr{testAddr}, nil,
+	).Once()
 
 	// Next, we'll make a link node with the same pubkey, but with an
 	// additional address.
-	nodePub, err := testNode.PubKey()
-	require.NoError(t, err, "unable to recv node pub")
 	linkNode := NewLinkNode(
 		fullDB.channelStateDB.linkNodeDB, wire.MainNet, nodePub,
 		anotherAddr,
diff --git a/channeldb/meta_test.go b/channeldb/meta_test.go
index 5b6bd29a9..97f6f0489 100644
--- a/channeldb/meta_test.go
+++ b/channeldb/meta_test.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/btcsuite/btcwallet/walletdb"
 	"github.com/go-errors/errors"
-	graphdb "github.com/lightningnetwork/lnd/graph/db"
 	"github.com/lightningnetwork/lnd/kvdb"
 	"github.com/stretchr/testify/require"
 )
@@ -22,13 +21,6 @@ func applyMigration(t *testing.T, beforeMigration, afterMigration func(d *DB),
 	}
 	cdb.dryRun = dryRun
 
-	// Create a test node that will be our source node.
-	testNode := createTestVertex(t)
-
-	graph, err := graphdb.MakeTestGraph(t)
-	require.NoError(t, err)
-	require.NoError(t, graph.SetSourceNode(testNode))
-
 	// beforeMigration usually used for populating the database
 	// with test data.
 	beforeMigration(cdb)