autopilot/prefattach_test: use NodeScores API for prefAttach tests

This commit converts the existing unit tests of Select into tests of
NodeScores.
This commit is contained in:
Johan T. Halseth 2018-11-22 23:18:09 +01:00
parent 5e8e54083f
commit be45697c6d
No known key found for this signature in database
GPG Key ID: 15BAADA29DA20D26

@ -230,10 +230,8 @@ var chanGraphs = []struct {
}, },
} }
// TestConstrainedPrefAttachmentSelectEmptyGraph ensures that when passed en // TestConstrainedPrefAttachmentSelectEmptyGraph ensures that when passed an
// empty graph, the Select function always detects the state, and returns nil. // empty graph, the NodeSores function always returns a score of 0.
// Otherwise, it would be possible for the main Select loop to entire an
// infinite loop.
func TestConstrainedPrefAttachmentSelectEmptyGraph(t *testing.T) { func TestConstrainedPrefAttachmentSelectEmptyGraph(t *testing.T) {
const ( const (
minChanSize = 0 minChanSize = 0
@ -249,15 +247,18 @@ func TestConstrainedPrefAttachmentSelectEmptyGraph(t *testing.T) {
Allocation: threshold, Allocation: threshold,
} }
// First, we'll generate a random key that represents "us", and create
// a new instance of the heuristic with our set parameters.
self, err := randKey()
if err != nil {
t.Fatalf("unable to generate self key: %v", err)
}
prefAttach := NewConstrainedPrefAttachment(constraints) prefAttach := NewConstrainedPrefAttachment(constraints)
skipNodes := make(map[NodeID]struct{}) // Create a random public key, which we will query to get a score for.
pub, err := randKey()
if err != nil {
t.Fatalf("unable to generate key: %v", err)
}
nodes := map[NodeID]struct{}{
NewNodeID(pub): {},
}
for _, graph := range chanGraphs { for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) { success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc() graph, cleanup, err := graph.genFunc()
@ -268,23 +269,21 @@ func TestConstrainedPrefAttachmentSelectEmptyGraph(t *testing.T) {
defer cleanup() defer cleanup()
} }
// With the necessary state initialized, we'll not // With the necessary state initialized, we'll now
// attempt to select a set of candidates channel for // attempt to get the score for this one node.
// creation given the current state of the graph.
const walletFunds = btcutil.SatoshiPerBitcoin const walletFunds = btcutil.SatoshiPerBitcoin
directives, err := prefAttach.Select(self, graph, scores, err := prefAttach.NodeScores(graph, nil,
walletFunds, 5, skipNodes) walletFunds, nodes)
if err != nil { if err != nil {
t1.Fatalf("unable to select attachment "+ t1.Fatalf("unable to select attachment "+
"directives: %v", err) "directives: %v", err)
} }
// We shouldn't have selected any new directives as we // Since the graph is empty, we expect the score to be
// started with an empty graph. // 0, giving an empty return map.
if len(directives) != 0 { if len(scores) != 0 {
t1.Fatalf("zero attachment directives "+ t1.Fatalf("expected empty score map, "+
"should have been returned instead %v were", "instead got %v ", len(scores))
len(directives))
} }
}) })
if !success { if !success {
@ -293,9 +292,50 @@ func TestConstrainedPrefAttachmentSelectEmptyGraph(t *testing.T) {
} }
} }
// completeGraph is a helper method that adds numNodes fully connected nodes to
// the graph.
func completeGraph(t *testing.T, g testGraph, numNodes int) {
const chanCapacity = btcutil.SatoshiPerBitcoin
nodes := make(map[int]*btcec.PublicKey)
for i := 0; i < numNodes; i++ {
for j := i + 1; j < numNodes; j++ {
node1 := nodes[i]
node2 := nodes[j]
edge1, edge2, err := g.addRandChannel(
node1, node2, chanCapacity)
if err != nil {
t.Fatalf("unable to generate channel: %v", err)
}
if node1 == nil {
pubKeyBytes := edge1.Peer.PubKey()
nodes[i], err = btcec.ParsePubKey(
pubKeyBytes[:], btcec.S256(),
)
if err != nil {
t.Fatalf("unable to parse pubkey: %v",
err)
}
}
if node2 == nil {
pubKeyBytes := edge2.Peer.PubKey()
nodes[j], err = btcec.ParsePubKey(
pubKeyBytes[:], btcec.S256(),
)
if err != nil {
t.Fatalf("unable to parse pubkey: %v",
err)
}
}
}
}
}
// TestConstrainedPrefAttachmentSelectTwoVertexes ensures that when passed a // TestConstrainedPrefAttachmentSelectTwoVertexes ensures that when passed a
// graph with only two eligible vertexes, then both are selected (without any // graph with only two eligible vertexes, then both are given the same score,
// repeats), and the funds are appropriately allocated across each peer. // and the funds are appropriately allocated across each peer.
func TestConstrainedPrefAttachmentSelectTwoVertexes(t *testing.T) { func TestConstrainedPrefAttachmentSelectTwoVertexes(t *testing.T) {
t.Parallel() t.Parallel()
@ -314,7 +354,6 @@ func TestConstrainedPrefAttachmentSelectTwoVertexes(t *testing.T) {
ChanLimit: chanLimit, ChanLimit: chanLimit,
Allocation: threshold, Allocation: threshold,
} }
skipNodes := make(map[NodeID]struct{})
for _, graph := range chanGraphs { for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) { success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc() graph, cleanup, err := graph.genFunc()
@ -325,13 +364,6 @@ func TestConstrainedPrefAttachmentSelectTwoVertexes(t *testing.T) {
defer cleanup() defer cleanup()
} }
// First, we'll generate a random key that represents
// "us", and create a new instance of the heuristic
// with our set parameters.
self, err := randKey()
if err != nil {
t1.Fatalf("unable to generate self key: %v", err)
}
prefAttach := NewConstrainedPrefAttachment(constraints) prefAttach := NewConstrainedPrefAttachment(constraints)
// For this set, we'll load the memory graph with two // For this set, we'll load the memory graph with two
@ -342,43 +374,67 @@ func TestConstrainedPrefAttachmentSelectTwoVertexes(t *testing.T) {
t1.Fatalf("unable to generate channel: %v", err) t1.Fatalf("unable to generate channel: %v", err)
} }
// With the necessary state initialized, we'll not // Get the score for all nodes found in the graph at
// attempt to select a set of candidates channel for // this point.
// creation given the current state of the graph. nodes := make(map[NodeID]struct{})
if err := graph.ForEachNode(func(n Node) error {
nodes[n.PubKey()] = struct{}{}
return nil
}); err != nil {
t1.Fatalf("unable to traverse graph: %v", err)
}
if len(nodes) != 2 {
t1.Fatalf("expected 2 nodes, found %d", len(nodes))
}
// With the necessary state initialized, we'll now
// attempt to get our candidates channel score given
// the current state of the graph.
const walletFunds = btcutil.SatoshiPerBitcoin * 10 const walletFunds = btcutil.SatoshiPerBitcoin * 10
directives, err := prefAttach.Select(self, graph, candidates, err := prefAttach.NodeScores(graph, nil,
walletFunds, 2, skipNodes) walletFunds, nodes)
if err != nil { if err != nil {
t1.Fatalf("unable to select attachment directives: %v", err) t1.Fatalf("unable to select attachment "+
"directives: %v", err)
} }
// Two new directives should have been selected, one if len(candidates) != len(nodes) {
// for each node already present within the graph. t1.Fatalf("all nodes should be scored, "+
if len(directives) != 2 { "instead %v were", len(candidates))
t1.Fatalf("two attachment directives should have been "+
"returned instead %v were", len(directives))
} }
// The node attached to should be amongst the two edges // The candidates should be amongst the two edges
// created above. // created above.
for _, directive := range directives { for nodeID, candidate := range candidates {
edge1Pub := edge1.Peer.PubKey() edge1Pub := edge1.Peer.PubKey()
edge2Pub := edge2.Peer.PubKey() edge2Pub := edge2.Peer.PubKey()
switch { switch {
case bytes.Equal(directive.NodeID[:], edge1Pub[:]): case bytes.Equal(nodeID[:], edge1Pub[:]):
case bytes.Equal(directive.NodeID[:], edge2Pub[:]): case bytes.Equal(nodeID[:], edge2Pub[:]):
default: default:
t1.Fatalf("attached to unknown node: %x", t1.Fatalf("attached to unknown node: %x",
directive.NodeID[:]) nodeID[:])
} }
// As the number of funds available exceed the // As the number of funds available exceed the
// max channel size, both edges should consume // max channel size, both edges should consume
// the maximum channel size. // the maximum channel size.
if directive.ChanAmt != maxChanSize { if candidate.ChanAmt != maxChanSize {
t1.Fatalf("max channel size should be allocated, "+ t1.Fatalf("max channel size should be "+
"instead %v was: ", maxChanSize) "allocated, instead %v was: ",
maxChanSize)
}
// Since each of the nodes has 1 channel, out
// of only one channel in the graph, we expect
// their score to be 0.5.
expScore := float64(0.5)
if candidate.Score != expScore {
t1.Fatalf("expected candidate score "+
"to be %v, instead was %v",
expScore, candidate.Score)
} }
} }
}) })
@ -410,7 +466,6 @@ func TestConstrainedPrefAttachmentSelectInsufficientFunds(t *testing.T) {
Allocation: threshold, Allocation: threshold,
} }
skipNodes := make(map[NodeID]struct{})
for _, graph := range chanGraphs { for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) { success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc() graph, cleanup, err := graph.genFunc()
@ -421,28 +476,36 @@ func TestConstrainedPrefAttachmentSelectInsufficientFunds(t *testing.T) {
defer cleanup() defer cleanup()
} }
// First, we'll generate a random key that represents // Add 10 nodes to the graph, with channels between
// "us", and create a new instance of the heuristic // them.
// with our set parameters. completeGraph(t, graph, 10)
self, err := randKey()
if err != nil {
t1.Fatalf("unable to generate self key: %v", err)
}
prefAttach := NewConstrainedPrefAttachment(constraints) prefAttach := NewConstrainedPrefAttachment(constraints)
// Next, we'll attempt to select a set of candidates, nodes := make(map[NodeID]struct{})
if err := graph.ForEachNode(func(n Node) error {
nodes[n.PubKey()] = struct{}{}
return nil
}); err != nil {
t1.Fatalf("unable to traverse graph: %v", err)
}
// With the necessary state initialized, we'll now
// attempt to get the score for our list of nodes,
// passing zero for the amount of wallet funds. This // passing zero for the amount of wallet funds. This
// should return an empty slice of directives. // should return an all-zero score set.
directives, err := prefAttach.Select(self, graph, 0, scores, err := prefAttach.NodeScores(graph, nil,
0, skipNodes) 0, nodes)
if err != nil { if err != nil {
t1.Fatalf("unable to select attachment "+ t1.Fatalf("unable to select attachment "+
"directives: %v", err) "directives: %v", err)
} }
if len(directives) != 0 {
t1.Fatalf("zero attachment directives "+ // Since all should be given a score of 0, the map
"should have been returned instead %v were", // should be empty.
len(directives)) if len(scores) != 0 {
t1.Fatalf("expected empty score map, "+
"instead got %v ", len(scores))
} }
}) })
if !success { if !success {
@ -452,9 +515,8 @@ func TestConstrainedPrefAttachmentSelectInsufficientFunds(t *testing.T) {
} }
// TestConstrainedPrefAttachmentSelectGreedyAllocation tests that if upon // TestConstrainedPrefAttachmentSelectGreedyAllocation tests that if upon
// deciding a set of candidates, we're unable to evenly split our funds, then // returning node scores, the NodeScores method will attempt to greedily
// we attempt to greedily allocate all funds to each selected vertex (up to the // allocate all funds to each vertex (up to the max channel size).
// max channel size).
func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) { func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) {
t.Parallel() t.Parallel()
@ -474,7 +536,6 @@ func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) {
Allocation: threshold, Allocation: threshold,
} }
skipNodes := make(map[NodeID]struct{})
for _, graph := range chanGraphs { for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) { success := t.Run(graph.name, func(t1 *testing.T) {
graph, cleanup, err := graph.genFunc() graph, cleanup, err := graph.genFunc()
@ -485,13 +546,6 @@ func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) {
defer cleanup() defer cleanup()
} }
// First, we'll generate a random key that represents
// "us", and create a new instance of the heuristic
// with our set parameters.
self, err := randKey()
if err != nil {
t1.Fatalf("unable to generate self key: %v", err)
}
prefAttach := NewConstrainedPrefAttachment(constraints) prefAttach := NewConstrainedPrefAttachment(constraints)
const chanCapacity = btcutil.SatoshiPerBitcoin const chanCapacity = btcutil.SatoshiPerBitcoin
@ -521,9 +575,10 @@ func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) {
// graph, with node node having two edges. // graph, with node node having two edges.
numNodes := 0 numNodes := 0
twoChans := false twoChans := false
nodes := make(map[NodeID]struct{})
if err := graph.ForEachNode(func(n Node) error { if err := graph.ForEachNode(func(n Node) error {
numNodes++ numNodes++
nodes[n.PubKey()] = struct{}{}
numChans := 0 numChans := 0
err := n.ForEachChannel(func(c ChannelEdge) error { err := n.ForEachChannel(func(c ChannelEdge) error {
numChans++ numChans++
@ -553,38 +608,61 @@ func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) {
// result, the heuristic should try to greedily // result, the heuristic should try to greedily
// allocate funds to channels. // allocate funds to channels.
const availableBalance = btcutil.SatoshiPerBitcoin * 2.5 const availableBalance = btcutil.SatoshiPerBitcoin * 2.5
directives, err := prefAttach.Select(self, graph, scores, err := prefAttach.NodeScores(graph, nil,
availableBalance, 5, skipNodes) availableBalance, nodes)
if err != nil { if err != nil {
t1.Fatalf("unable to select attachment "+ t1.Fatalf("unable to select attachment "+
"directives: %v", err) "directives: %v", err)
} }
// Three directives should have been returned. if len(scores) != len(nodes) {
if len(directives) != 3 { t1.Fatalf("all nodes should be scored, "+
t1.Fatalf("expected 3 directives, instead "+ "instead %v were", len(scores))
"got: %v", len(directives))
} }
// The two directive should have the max channel amount // The candidates should have a non-zero score, and
// allocated. // have the max chan size funds recommended channel
if directives[0].ChanAmt != maxChanSize { // size.
t1.Fatalf("expected recommendation of %v, "+ for _, candidate := range scores {
"instead got %v", maxChanSize, if candidate.Score == 0 {
directives[0].ChanAmt) t1.Fatalf("Expected non-zero score")
}
if directives[1].ChanAmt != maxChanSize {
t1.Fatalf("expected recommendation of %v, "+
"instead got %v", maxChanSize,
directives[1].ChanAmt)
} }
// The third channel should have been allocated the if candidate.ChanAmt != maxChanSize {
// remainder, or 0.5 BTC. t1.Fatalf("expected recommendation "+
if directives[2].ChanAmt != (btcutil.SatoshiPerBitcoin * 0.5) { "of %v, instead got %v",
t1.Fatalf("expected recommendation of %v, "+ maxChanSize, candidate.ChanAmt)
"instead got %v", maxChanSize, }
directives[2].ChanAmt) }
// Imagine a few channels are being opened, and there's
// only 0.5 BTC left. That should leave us with channel
// candidates of that size.
const remBalance = btcutil.SatoshiPerBitcoin * 0.5
scores, err = prefAttach.NodeScores(graph, nil,
remBalance, nodes)
if err != nil {
t1.Fatalf("unable to select attachment "+
"directives: %v", err)
}
if len(scores) != len(nodes) {
t1.Fatalf("all nodes should be scored, "+
"instead %v were", len(scores))
}
// Check that the recommended channel sizes are now the
// remaining channel balance.
for _, candidate := range scores {
if candidate.Score == 0 {
t1.Fatalf("Expected non-zero score")
}
if candidate.ChanAmt != remBalance {
t1.Fatalf("expected recommendation "+
"of %v, instead got %v",
remBalance, candidate.ChanAmt)
}
} }
}) })
if !success { if !success {
@ -594,8 +672,8 @@ func TestConstrainedPrefAttachmentSelectGreedyAllocation(t *testing.T) {
} }
// TestConstrainedPrefAttachmentSelectSkipNodes ensures that if a node was // TestConstrainedPrefAttachmentSelectSkipNodes ensures that if a node was
// already select for attachment, then that node is excluded from the set of // already selected as a channel counterparty, then that node will get a score
// candidate nodes. // of zero during scoring.
func TestConstrainedPrefAttachmentSelectSkipNodes(t *testing.T) { func TestConstrainedPrefAttachmentSelectSkipNodes(t *testing.T) {
t.Parallel() t.Parallel()
@ -617,8 +695,6 @@ func TestConstrainedPrefAttachmentSelectSkipNodes(t *testing.T) {
for _, graph := range chanGraphs { for _, graph := range chanGraphs {
success := t.Run(graph.name, func(t1 *testing.T) { success := t.Run(graph.name, func(t1 *testing.T) {
skipNodes := make(map[NodeID]struct{})
graph, cleanup, err := graph.genFunc() graph, cleanup, err := graph.genFunc()
if err != nil { if err != nil {
t1.Fatalf("unable to create graph: %v", err) t1.Fatalf("unable to create graph: %v", err)
@ -627,13 +703,6 @@ func TestConstrainedPrefAttachmentSelectSkipNodes(t *testing.T) {
defer cleanup() defer cleanup()
} }
// First, we'll generate a random key that represents
// "us", and create a new instance of the heuristic
// with our set parameters.
self, err := randKey()
if err != nil {
t1.Fatalf("unable to generate self key: %v", err)
}
prefAttach := NewConstrainedPrefAttachment(constraints) prefAttach := NewConstrainedPrefAttachment(constraints)
// Next, we'll create a simple topology of two nodes, // Next, we'll create a simple topology of two nodes,
@ -645,44 +714,74 @@ func TestConstrainedPrefAttachmentSelectSkipNodes(t *testing.T) {
t1.Fatalf("unable to create channel: %v", err) t1.Fatalf("unable to create channel: %v", err)
} }
// With our graph created, we'll now execute the Select nodes := make(map[NodeID]struct{})
// function to recommend potential attachment if err := graph.ForEachNode(func(n Node) error {
// candidates. nodes[n.PubKey()] = struct{}{}
return nil
}); err != nil {
t1.Fatalf("unable to traverse graph: %v", err)
}
if len(nodes) != 2 {
t1.Fatalf("expected 2 nodes, found %d", len(nodes))
}
// With our graph created, we'll now get the scores for
// all nodes in the graph.
const availableBalance = btcutil.SatoshiPerBitcoin * 2.5 const availableBalance = btcutil.SatoshiPerBitcoin * 2.5
directives, err := prefAttach.Select(self, graph, scores, err := prefAttach.NodeScores(graph, nil,
availableBalance, 2, skipNodes) availableBalance, nodes)
if err != nil { if err != nil {
t1.Fatalf("unable to select attachment "+ t1.Fatalf("unable to select attachment "+
"directives: %v", err) "directives: %v", err)
} }
// As the channel limit is three, and two nodes are if len(scores) != len(nodes) {
// present in the graph, both should be selected. t1.Fatalf("all nodes should be scored, "+
if len(directives) != 2 { "instead %v were", len(scores))
t1.Fatalf("expected two directives, instead "+ }
"got %v", len(directives))
// THey should all have a score, and a maxChanSize
// channel size recommendation.
for _, candidate := range scores {
if candidate.Score == 0 {
t1.Fatalf("Expected non-zero score")
}
if candidate.ChanAmt != maxChanSize {
t1.Fatalf("expected recommendation "+
"of %v, instead got %v",
maxChanSize, candidate.ChanAmt)
}
} }
// We'll simulate a channel update by adding the nodes // We'll simulate a channel update by adding the nodes
// we just establish channel with the to set of nodes // to our set of channels.
// to be skipped. var chans []Channel
skipNodes[directives[0].NodeID] = struct{}{} for _, candidate := range scores {
skipNodes[directives[1].NodeID] = struct{}{} chans = append(chans,
Channel{
Node: candidate.NodeID,
},
)
}
// If we attempt to make a call to the Select function, // If we attempt to make a call to the NodeScores
// without providing any new information, then we // function, without providing any new information,
// should get no new directives as both nodes has // then all nodes should have a score of zero, since we
// already been attached to. // already got channels to them.
directives, err = prefAttach.Select(self, graph, scores, err = prefAttach.NodeScores(graph, chans,
availableBalance, 2, skipNodes) availableBalance, nodes)
if err != nil { if err != nil {
t1.Fatalf("unable to select attachment "+ t1.Fatalf("unable to select attachment "+
"directives: %v", err) "directives: %v", err)
} }
if len(directives) != 0 { // Since all should be given a score of 0, the map
t1.Fatalf("zero new directives should have been "+ // should be empty.
"selected, but %v were", len(directives)) if len(scores) != 0 {
t1.Fatalf("expected empty score map, "+
"instead got %v ", len(scores))
} }
}) })
if !success { if !success {