lnd.xprv/watchtower/wtclient/candidate_iterator_test.go

158 lines
4.5 KiB
Go
Raw Normal View History

package wtclient
import (
"encoding/binary"
"math/rand"
"net"
"reflect"
"testing"
"time"
"github.com/btcsuite/btcd/btcec"
"github.com/davecgh/go-spew/spew"
"github.com/lightningnetwork/lnd/watchtower/wtdb"
)
func init() {
rand.Seed(time.Now().Unix())
}
func randAddr(t *testing.T) net.Addr {
var ip [4]byte
if _, err := rand.Read(ip[:]); err != nil {
t.Fatal(err)
}
var port [2]byte
if _, err := rand.Read(port[:]); err != nil {
t.Fatal(err)
}
return &net.TCPAddr{
IP: net.IP(ip[:]),
Port: int(binary.BigEndian.Uint16(port[:])),
}
}
func randTower(t *testing.T) *wtdb.Tower {
priv, err := btcec.NewPrivateKey(btcec.S256())
if err != nil {
t.Fatalf("unable to create private key: %v", err)
}
pubKey := priv.PubKey()
pubKey.Curve = nil
return &wtdb.Tower{
ID: wtdb.TowerID(rand.Uint64()),
IdentityKey: pubKey,
Addresses: []net.Addr{randAddr(t)},
}
}
func copyTower(tower *wtdb.Tower) *wtdb.Tower {
t := &wtdb.Tower{
ID: tower.ID,
IdentityKey: tower.IdentityKey,
Addresses: make([]net.Addr, len(tower.Addresses)),
}
copy(t.Addresses, tower.Addresses)
return t
}
func assertActiveCandidate(t *testing.T, i TowerCandidateIterator,
c *wtdb.Tower, active bool) {
isCandidate := i.IsActive(c.ID)
if isCandidate && !active {
t.Fatalf("expected tower %v to no longer be an active candidate",
c.ID)
}
if !isCandidate && active {
t.Fatalf("expected tower %v to be an active candidate", c.ID)
}
}
func assertNextCandidate(t *testing.T, i TowerCandidateIterator, c *wtdb.Tower) {
t.Helper()
tower, err := i.Next()
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(tower, c) {
t.Fatalf("expected tower: %v\ngot: %v", spew.Sdump(c),
spew.Sdump(tower))
}
}
// TestTowerCandidateIterator asserts the internal state of a
// TowerCandidateIterator after a series of updates to its candidates.
func TestTowerCandidateIterator(t *testing.T) {
t.Parallel()
// We'll start our test by creating an iterator of four candidate
// towers. We'll use copies of these towers within the iterator to
// ensure the iterator properly updates the state of its candidates.
const numTowers = 4
towers := make([]*wtdb.Tower, 0, numTowers)
for i := 0; i < numTowers; i++ {
towers = append(towers, randTower(t))
}
towerCopies := make([]*wtdb.Tower, 0, numTowers)
for _, tower := range towers {
towerCopies = append(towerCopies, copyTower(tower))
}
towerIterator := newTowerListIterator(towerCopies...)
// We should expect to see all of our candidates in the order that they
// were added.
for _, expTower := range towers {
tower, err := towerIterator.Next()
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(tower, expTower) {
t.Fatalf("expected tower: %v\ngot: %v",
spew.Sdump(expTower), spew.Sdump(tower))
}
}
if _, err := towerIterator.Next(); err != ErrTowerCandidatesExhausted {
t.Fatalf("expected ErrTowerCandidatesExhausted, got %v", err)
}
towerIterator.Reset()
// We'll then attempt to test the RemoveCandidate behavior of the
// iterator. We'll remove the address of the first tower, which should
// result in it not having any addresses left, but still being an active
// candidate.
firstTower := towers[0]
firstTowerAddr := firstTower.Addresses[0]
firstTower.RemoveAddress(firstTowerAddr)
towerIterator.RemoveCandidate(firstTower.ID, firstTowerAddr)
assertActiveCandidate(t, towerIterator, firstTower, true)
assertNextCandidate(t, towerIterator, firstTower)
// We'll then remove the second tower completely from the iterator by
// not providing the optional address. Since it's been removed, we
// should expect to see the third tower next.
secondTower, thirdTower := towers[1], towers[2]
towerIterator.RemoveCandidate(secondTower.ID, nil)
assertActiveCandidate(t, towerIterator, secondTower, false)
assertNextCandidate(t, towerIterator, thirdTower)
// We'll then update the fourth candidate with a new address. A
// duplicate shouldn't be added since it already exists within the
// iterator, but the new address should be.
fourthTower := towers[3]
assertActiveCandidate(t, towerIterator, fourthTower, true)
fourthTower.AddAddress(randAddr(t))
towerIterator.AddCandidate(fourthTower)
assertNextCandidate(t, towerIterator, fourthTower)
// Finally, we'll attempt to add a new candidate to the end of the
// iterator. Since it didn't already exist and we've reached the end, it
// should be available as the next candidate.
towerIterator.AddCandidate(secondTower)
assertActiveCandidate(t, towerIterator, secondTower, true)
assertNextCandidate(t, towerIterator, secondTower)
}