From 1184c9eaf8df1b4d222b02e47208d51c17c98383 Mon Sep 17 00:00:00 2001 From: Elle Mouton Date: Wed, 5 Feb 2025 12:11:28 +0200 Subject: [PATCH] autopilot: move tests code to test files This is a pure code move commit where we move any code that is only ever used by tests to test files. Many of the calls to the graphdb.ChannelGraph pointer are only coming from tests code. --- autopilot/graph.go | 328 ----------------------------------- autopilot/prefattach_test.go | 326 ++++++++++++++++++++++++++++++++++ 2 files changed, 326 insertions(+), 328 deletions(-) diff --git a/autopilot/graph.go b/autopilot/graph.go index b4e415077..b14d48231 100644 --- a/autopilot/graph.go +++ b/autopilot/graph.go @@ -1,13 +1,9 @@ package autopilot import ( - "bytes" "encoding/hex" - "errors" "net" "sort" - "sync/atomic" - "time" "github.com/btcsuite/btcd/btcec/v2" "github.com/btcsuite/btcd/btcec/v2/ecdsa" @@ -154,330 +150,6 @@ func (d *databaseChannelGraph) ForEachNode(cb func(Node) error) error { }) } -// addRandChannel creates a new channel two target nodes. This function is -// meant to aide in the generation of random graphs for use within test cases -// the exercise the autopilot package. -func (d *databaseChannelGraph) addRandChannel(node1, node2 *btcec.PublicKey, - capacity btcutil.Amount) (*ChannelEdge, *ChannelEdge, error) { - - fetchNode := func(pub *btcec.PublicKey) (*models.LightningNode, error) { - if pub != nil { - vertex, err := route.NewVertexFromBytes( - pub.SerializeCompressed(), - ) - if err != nil { - return nil, err - } - - dbNode, err := d.db.FetchLightningNode(vertex) - switch { - case errors.Is(err, graphdb.ErrGraphNodeNotFound): - fallthrough - case errors.Is(err, graphdb.ErrGraphNotFound): - graphNode := &models.LightningNode{ - HaveNodeAnnouncement: true, - Addresses: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - Features: lnwire.NewFeatureVector( - nil, lnwire.Features, - ), - AuthSigBytes: testSig.Serialize(), - } - graphNode.AddPubKey(pub) - if err := d.db.AddLightningNode(graphNode); err != nil { - return nil, err - } - case err != nil: - return nil, err - } - - return dbNode, nil - } - - nodeKey, err := randKey() - if err != nil { - return nil, err - } - dbNode := &models.LightningNode{ - HaveNodeAnnouncement: true, - Addresses: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - Features: lnwire.NewFeatureVector( - nil, lnwire.Features, - ), - AuthSigBytes: testSig.Serialize(), - } - dbNode.AddPubKey(nodeKey) - if err := d.db.AddLightningNode(dbNode); err != nil { - return nil, err - } - - return dbNode, nil - } - - vertex1, err := fetchNode(node1) - if err != nil { - return nil, nil, err - } - - vertex2, err := fetchNode(node2) - if err != nil { - return nil, nil, err - } - - var lnNode1, lnNode2 *btcec.PublicKey - if bytes.Compare(vertex1.PubKeyBytes[:], vertex2.PubKeyBytes[:]) == -1 { - lnNode1, _ = vertex1.PubKey() - lnNode2, _ = vertex2.PubKey() - } else { - lnNode1, _ = vertex2.PubKey() - lnNode2, _ = vertex1.PubKey() - } - - chanID := randChanID() - edge := &models.ChannelEdgeInfo{ - ChannelID: chanID.ToUint64(), - Capacity: capacity, - } - edge.AddNodeKeys(lnNode1, lnNode2, lnNode1, lnNode2) - if err := d.db.AddChannelEdge(edge); err != nil { - return nil, nil, err - } - edgePolicy := &models.ChannelEdgePolicy{ - SigBytes: testSig.Serialize(), - ChannelID: chanID.ToUint64(), - LastUpdate: time.Now(), - TimeLockDelta: 10, - MinHTLC: 1, - MaxHTLC: lnwire.NewMSatFromSatoshis(capacity), - FeeBaseMSat: 10, - FeeProportionalMillionths: 10000, - MessageFlags: 1, - ChannelFlags: 0, - } - - if err := d.db.UpdateEdgePolicy(edgePolicy); err != nil { - return nil, nil, err - } - edgePolicy = &models.ChannelEdgePolicy{ - SigBytes: testSig.Serialize(), - ChannelID: chanID.ToUint64(), - LastUpdate: time.Now(), - TimeLockDelta: 10, - MinHTLC: 1, - MaxHTLC: lnwire.NewMSatFromSatoshis(capacity), - FeeBaseMSat: 10, - FeeProportionalMillionths: 10000, - MessageFlags: 1, - ChannelFlags: 1, - } - if err := d.db.UpdateEdgePolicy(edgePolicy); err != nil { - return nil, nil, err - } - - return &ChannelEdge{ - ChanID: chanID, - Capacity: capacity, - Peer: &dbNode{ - db: d.db, - node: vertex1, - }, - }, - &ChannelEdge{ - ChanID: chanID, - Capacity: capacity, - Peer: &dbNode{ - db: d.db, - node: vertex2, - }, - }, - nil -} - -func (d *databaseChannelGraph) addRandNode() (*btcec.PublicKey, error) { - nodeKey, err := randKey() - if err != nil { - return nil, err - } - dbNode := &models.LightningNode{ - HaveNodeAnnouncement: true, - Addresses: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - Features: lnwire.NewFeatureVector( - nil, lnwire.Features, - ), - AuthSigBytes: testSig.Serialize(), - } - dbNode.AddPubKey(nodeKey) - if err := d.db.AddLightningNode(dbNode); err != nil { - return nil, err - } - - return nodeKey, nil - -} - -// memChannelGraph is an implementation of the autopilot.ChannelGraph backed by -// an in-memory graph. -type memChannelGraph struct { - graph map[NodeID]*memNode -} - -// A compile time assertion to ensure memChannelGraph meets the -// autopilot.ChannelGraph interface. -var _ ChannelGraph = (*memChannelGraph)(nil) - -// newMemChannelGraph creates a new blank in-memory channel graph -// implementation. -func newMemChannelGraph() *memChannelGraph { - return &memChannelGraph{ - graph: make(map[NodeID]*memNode), - } -} - -// ForEachNode is a higher-order function that should be called once for each -// connected node within the channel graph. If the passed callback returns an -// error, then execution should be terminated. -// -// NOTE: Part of the autopilot.ChannelGraph interface. -func (m memChannelGraph) ForEachNode(cb func(Node) error) error { - for _, node := range m.graph { - if err := cb(node); err != nil { - return err - } - } - - return nil -} - -// randChanID generates a new random channel ID. -func randChanID() lnwire.ShortChannelID { - id := atomic.AddUint64(&chanIDCounter, 1) - return lnwire.NewShortChanIDFromInt(id) -} - -// randKey returns a random public key. -func randKey() (*btcec.PublicKey, error) { - priv, err := btcec.NewPrivateKey() - if err != nil { - return nil, err - } - - return priv.PubKey(), nil -} - -// addRandChannel creates a new channel two target nodes. This function is -// meant to aide in the generation of random graphs for use within test cases -// the exercise the autopilot package. -func (m *memChannelGraph) addRandChannel(node1, node2 *btcec.PublicKey, - capacity btcutil.Amount) (*ChannelEdge, *ChannelEdge, error) { - - var ( - vertex1, vertex2 *memNode - ok bool - ) - - if node1 != nil { - vertex1, ok = m.graph[NewNodeID(node1)] - if !ok { - vertex1 = &memNode{ - pub: node1, - addrs: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - } - } - } else { - newPub, err := randKey() - if err != nil { - return nil, nil, err - } - vertex1 = &memNode{ - pub: newPub, - addrs: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - } - } - - if node2 != nil { - vertex2, ok = m.graph[NewNodeID(node2)] - if !ok { - vertex2 = &memNode{ - pub: node2, - addrs: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - } - } - } else { - newPub, err := randKey() - if err != nil { - return nil, nil, err - } - vertex2 = &memNode{ - pub: newPub, - addrs: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - } - } - - edge1 := ChannelEdge{ - ChanID: randChanID(), - Capacity: capacity, - Peer: vertex2, - } - vertex1.chans = append(vertex1.chans, edge1) - - edge2 := ChannelEdge{ - ChanID: randChanID(), - Capacity: capacity, - Peer: vertex1, - } - vertex2.chans = append(vertex2.chans, edge2) - - m.graph[NewNodeID(vertex1.pub)] = vertex1 - m.graph[NewNodeID(vertex2.pub)] = vertex2 - - return &edge1, &edge2, nil -} - -func (m *memChannelGraph) addRandNode() (*btcec.PublicKey, error) { - newPub, err := randKey() - if err != nil { - return nil, err - } - vertex := &memNode{ - pub: newPub, - addrs: []net.Addr{ - &net.TCPAddr{ - IP: bytes.Repeat([]byte("a"), 16), - }, - }, - } - m.graph[NewNodeID(newPub)] = vertex - - return newPub, nil -} - // databaseChannelGraphCached wraps a channeldb.ChannelGraph instance with the // necessary API to properly implement the autopilot.ChannelGraph interface. type databaseChannelGraphCached struct { diff --git a/autopilot/prefattach_test.go b/autopilot/prefattach_test.go index ab52c55f6..37fbdaa53 100644 --- a/autopilot/prefattach_test.go +++ b/autopilot/prefattach_test.go @@ -2,14 +2,20 @@ package autopilot import ( "bytes" + "errors" prand "math/rand" + "net" + "sync/atomic" "testing" "time" "github.com/btcsuite/btcd/btcec/v2" "github.com/btcsuite/btcd/btcutil" graphdb "github.com/lightningnetwork/lnd/graph/db" + "github.com/lightningnetwork/lnd/graph/db/models" "github.com/lightningnetwork/lnd/kvdb" + "github.com/lightningnetwork/lnd/lnwire" + "github.com/lightningnetwork/lnd/routing/route" "github.com/stretchr/testify/require" ) @@ -368,3 +374,323 @@ func TestPrefAttachmentSelectSkipNodes(t *testing.T) { } } } + +// addRandChannel creates a new channel two target nodes. This function is +// meant to aide in the generation of random graphs for use within test cases +// the exercise the autopilot package. +func (d *databaseChannelGraph) addRandChannel(node1, node2 *btcec.PublicKey, + capacity btcutil.Amount) (*ChannelEdge, *ChannelEdge, error) { + + fetchNode := func(pub *btcec.PublicKey) (*models.LightningNode, error) { + if pub != nil { + vertex, err := route.NewVertexFromBytes( + pub.SerializeCompressed(), + ) + if err != nil { + return nil, err + } + + dbNode, err := d.db.FetchLightningNode(vertex) + switch { + case errors.Is(err, graphdb.ErrGraphNodeNotFound): + fallthrough + case errors.Is(err, graphdb.ErrGraphNotFound): + graphNode := &models.LightningNode{ + HaveNodeAnnouncement: true, + Addresses: []net.Addr{&net.TCPAddr{ + IP: bytes.Repeat( + []byte("a"), 16, + ), + }}, + Features: lnwire.NewFeatureVector( + nil, lnwire.Features, + ), + AuthSigBytes: testSig.Serialize(), + } + graphNode.AddPubKey(pub) + err := d.db.AddLightningNode(graphNode) + if err != nil { + return nil, err + } + case err != nil: + return nil, err + } + + return dbNode, nil + } + + nodeKey, err := randKey() + if err != nil { + return nil, err + } + dbNode := &models.LightningNode{ + HaveNodeAnnouncement: true, + Addresses: []net.Addr{ + &net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }, + }, + Features: lnwire.NewFeatureVector( + nil, lnwire.Features, + ), + AuthSigBytes: testSig.Serialize(), + } + dbNode.AddPubKey(nodeKey) + if err := d.db.AddLightningNode(dbNode); err != nil { + return nil, err + } + + return dbNode, nil + } + + vertex1, err := fetchNode(node1) + if err != nil { + return nil, nil, err + } + + vertex2, err := fetchNode(node2) + if err != nil { + return nil, nil, err + } + + var lnNode1, lnNode2 *btcec.PublicKey + if bytes.Compare(vertex1.PubKeyBytes[:], vertex2.PubKeyBytes[:]) == -1 { + lnNode1, _ = vertex1.PubKey() + lnNode2, _ = vertex2.PubKey() + } else { + lnNode1, _ = vertex2.PubKey() + lnNode2, _ = vertex1.PubKey() + } + + chanID := randChanID() + edge := &models.ChannelEdgeInfo{ + ChannelID: chanID.ToUint64(), + Capacity: capacity, + } + edge.AddNodeKeys(lnNode1, lnNode2, lnNode1, lnNode2) + if err := d.db.AddChannelEdge(edge); err != nil { + return nil, nil, err + } + edgePolicy := &models.ChannelEdgePolicy{ + SigBytes: testSig.Serialize(), + ChannelID: chanID.ToUint64(), + LastUpdate: time.Now(), + TimeLockDelta: 10, + MinHTLC: 1, + MaxHTLC: lnwire.NewMSatFromSatoshis(capacity), + FeeBaseMSat: 10, + FeeProportionalMillionths: 10000, + MessageFlags: 1, + ChannelFlags: 0, + } + + if err := d.db.UpdateEdgePolicy(edgePolicy); err != nil { + return nil, nil, err + } + edgePolicy = &models.ChannelEdgePolicy{ + SigBytes: testSig.Serialize(), + ChannelID: chanID.ToUint64(), + LastUpdate: time.Now(), + TimeLockDelta: 10, + MinHTLC: 1, + MaxHTLC: lnwire.NewMSatFromSatoshis(capacity), + FeeBaseMSat: 10, + FeeProportionalMillionths: 10000, + MessageFlags: 1, + ChannelFlags: 1, + } + if err := d.db.UpdateEdgePolicy(edgePolicy); err != nil { + return nil, nil, err + } + + return &ChannelEdge{ + ChanID: chanID, + Capacity: capacity, + Peer: &dbNode{ + db: d.db, + node: vertex1, + }, + }, + &ChannelEdge{ + ChanID: chanID, + Capacity: capacity, + Peer: &dbNode{ + db: d.db, + node: vertex2, + }, + }, + nil +} + +func (d *databaseChannelGraph) addRandNode() (*btcec.PublicKey, error) { + nodeKey, err := randKey() + if err != nil { + return nil, err + } + dbNode := &models.LightningNode{ + HaveNodeAnnouncement: true, + Addresses: []net.Addr{ + &net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }, + }, + Features: lnwire.NewFeatureVector( + nil, lnwire.Features, + ), + AuthSigBytes: testSig.Serialize(), + } + dbNode.AddPubKey(nodeKey) + if err := d.db.AddLightningNode(dbNode); err != nil { + return nil, err + } + + return nodeKey, nil +} + +// memChannelGraph is an implementation of the autopilot.ChannelGraph backed by +// an in-memory graph. +type memChannelGraph struct { + graph map[NodeID]*memNode +} + +// A compile time assertion to ensure memChannelGraph meets the +// autopilot.ChannelGraph interface. +var _ ChannelGraph = (*memChannelGraph)(nil) + +// newMemChannelGraph creates a new blank in-memory channel graph +// implementation. +func newMemChannelGraph() *memChannelGraph { + return &memChannelGraph{ + graph: make(map[NodeID]*memNode), + } +} + +// ForEachNode is a higher-order function that should be called once for each +// connected node within the channel graph. If the passed callback returns an +// error, then execution should be terminated. +// +// NOTE: Part of the autopilot.ChannelGraph interface. +func (m memChannelGraph) ForEachNode(cb func(Node) error) error { + for _, node := range m.graph { + if err := cb(node); err != nil { + return err + } + } + + return nil +} + +// randChanID generates a new random channel ID. +func randChanID() lnwire.ShortChannelID { + id := atomic.AddUint64(&chanIDCounter, 1) + return lnwire.NewShortChanIDFromInt(id) +} + +// randKey returns a random public key. +func randKey() (*btcec.PublicKey, error) { + priv, err := btcec.NewPrivateKey() + if err != nil { + return nil, err + } + + return priv.PubKey(), nil +} + +// addRandChannel creates a new channel two target nodes. This function is +// meant to aide in the generation of random graphs for use within test cases +// the exercise the autopilot package. +func (m *memChannelGraph) addRandChannel(node1, node2 *btcec.PublicKey, + capacity btcutil.Amount) (*ChannelEdge, *ChannelEdge, error) { + + var ( + vertex1, vertex2 *memNode + ok bool + ) + + if node1 != nil { + vertex1, ok = m.graph[NewNodeID(node1)] + if !ok { + vertex1 = &memNode{ + pub: node1, + addrs: []net.Addr{&net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }}, + } + } + } else { + newPub, err := randKey() + if err != nil { + return nil, nil, err + } + vertex1 = &memNode{ + pub: newPub, + addrs: []net.Addr{ + &net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }, + }, + } + } + + if node2 != nil { + vertex2, ok = m.graph[NewNodeID(node2)] + if !ok { + vertex2 = &memNode{ + pub: node2, + addrs: []net.Addr{&net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }}, + } + } + } else { + newPub, err := randKey() + if err != nil { + return nil, nil, err + } + vertex2 = &memNode{ + pub: newPub, + addrs: []net.Addr{ + &net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }, + }, + } + } + + edge1 := ChannelEdge{ + ChanID: randChanID(), + Capacity: capacity, + Peer: vertex2, + } + vertex1.chans = append(vertex1.chans, edge1) + + edge2 := ChannelEdge{ + ChanID: randChanID(), + Capacity: capacity, + Peer: vertex1, + } + vertex2.chans = append(vertex2.chans, edge2) + + m.graph[NewNodeID(vertex1.pub)] = vertex1 + m.graph[NewNodeID(vertex2.pub)] = vertex2 + + return &edge1, &edge2, nil +} + +func (m *memChannelGraph) addRandNode() (*btcec.PublicKey, error) { + newPub, err := randKey() + if err != nil { + return nil, err + } + vertex := &memNode{ + pub: newPub, + addrs: []net.Addr{ + &net.TCPAddr{ + IP: bytes.Repeat([]byte("a"), 16), + }, + }, + } + m.graph[NewNodeID(newPub)] = vertex + + return newPub, nil +}