github.com/decred/dcrlnd@v0.7.6/watchtower/wtclient/candidate_iterator_test.go (about)

     1  package wtclient
     2  
     3  import (
     4  	"encoding/binary"
     5  	"math/rand"
     6  	"net"
     7  	"reflect"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/davecgh/go-spew/spew"
    12  	"github.com/decred/dcrd/dcrec/secp256k1/v4"
    13  	"github.com/decred/dcrlnd/watchtower/wtdb"
    14  )
    15  
    16  func init() {
    17  	rand.Seed(time.Now().Unix())
    18  }
    19  
    20  func randAddr(t *testing.T) net.Addr {
    21  	var ip [4]byte
    22  	if _, err := rand.Read(ip[:]); err != nil {
    23  		t.Fatal(err)
    24  	}
    25  	var port [2]byte
    26  	if _, err := rand.Read(port[:]); err != nil {
    27  		t.Fatal(err)
    28  
    29  	}
    30  	return &net.TCPAddr{
    31  		IP:   net.IP(ip[:]),
    32  		Port: int(binary.BigEndian.Uint16(port[:])),
    33  	}
    34  }
    35  
    36  func randTower(t *testing.T) *wtdb.Tower {
    37  	priv, err := secp256k1.GeneratePrivateKey()
    38  	if err != nil {
    39  		t.Fatalf("unable to create private key: %v", err)
    40  	}
    41  	pubKey := priv.PubKey()
    42  	return &wtdb.Tower{
    43  		ID:          wtdb.TowerID(rand.Uint64()),
    44  		IdentityKey: pubKey,
    45  		Addresses:   []net.Addr{randAddr(t)},
    46  	}
    47  }
    48  
    49  func copyTower(tower *wtdb.Tower) *wtdb.Tower {
    50  	t := &wtdb.Tower{
    51  		ID:          tower.ID,
    52  		IdentityKey: tower.IdentityKey,
    53  		Addresses:   make([]net.Addr, len(tower.Addresses)),
    54  	}
    55  	copy(t.Addresses, tower.Addresses)
    56  	return t
    57  }
    58  
    59  func assertActiveCandidate(t *testing.T, i TowerCandidateIterator,
    60  	c *wtdb.Tower, active bool) {
    61  
    62  	isCandidate := i.IsActive(c.ID)
    63  	if isCandidate && !active {
    64  		t.Fatalf("expected tower %v to no longer be an active candidate",
    65  			c.ID)
    66  	}
    67  	if !isCandidate && active {
    68  		t.Fatalf("expected tower %v to be an active candidate", c.ID)
    69  	}
    70  }
    71  
    72  func assertNextCandidate(t *testing.T, i TowerCandidateIterator, c *wtdb.Tower) {
    73  	t.Helper()
    74  
    75  	tower, err := i.Next()
    76  	if err != nil {
    77  		t.Fatal(err)
    78  	}
    79  	if !reflect.DeepEqual(tower, c) {
    80  		t.Fatalf("expected tower: %v\ngot: %v", spew.Sdump(c),
    81  			spew.Sdump(tower))
    82  	}
    83  }
    84  
    85  // TestTowerCandidateIterator asserts the internal state of a
    86  // TowerCandidateIterator after a series of updates to its candidates.
    87  func TestTowerCandidateIterator(t *testing.T) {
    88  	t.Parallel()
    89  
    90  	// We'll start our test by creating an iterator of four candidate
    91  	// towers. We'll use copies of these towers within the iterator to
    92  	// ensure the iterator properly updates the state of its candidates.
    93  	const numTowers = 4
    94  	towers := make([]*wtdb.Tower, 0, numTowers)
    95  	for i := 0; i < numTowers; i++ {
    96  		towers = append(towers, randTower(t))
    97  	}
    98  	towerCopies := make([]*wtdb.Tower, 0, numTowers)
    99  	for _, tower := range towers {
   100  		towerCopies = append(towerCopies, copyTower(tower))
   101  	}
   102  	towerIterator := newTowerListIterator(towerCopies...)
   103  
   104  	// We should expect to see all of our candidates in the order that they
   105  	// were added.
   106  	for _, expTower := range towers {
   107  		tower, err := towerIterator.Next()
   108  		if err != nil {
   109  			t.Fatal(err)
   110  		}
   111  		if !reflect.DeepEqual(tower, expTower) {
   112  			t.Fatalf("expected tower: %v\ngot: %v",
   113  				spew.Sdump(expTower), spew.Sdump(tower))
   114  		}
   115  	}
   116  
   117  	if _, err := towerIterator.Next(); err != ErrTowerCandidatesExhausted {
   118  		t.Fatalf("expected ErrTowerCandidatesExhausted, got %v", err)
   119  	}
   120  	towerIterator.Reset()
   121  
   122  	// We'll then attempt to test the RemoveCandidate behavior of the
   123  	// iterator. We'll remove the address of the first tower, which should
   124  	// result in it not having any addresses left, but still being an active
   125  	// candidate.
   126  	firstTower := towers[0]
   127  	firstTowerAddr := firstTower.Addresses[0]
   128  	firstTower.RemoveAddress(firstTowerAddr)
   129  	towerIterator.RemoveCandidate(firstTower.ID, firstTowerAddr)
   130  	assertActiveCandidate(t, towerIterator, firstTower, true)
   131  	assertNextCandidate(t, towerIterator, firstTower)
   132  
   133  	// We'll then remove the second tower completely from the iterator by
   134  	// not providing the optional address. Since it's been removed, we
   135  	// should expect to see the third tower next.
   136  	secondTower, thirdTower := towers[1], towers[2]
   137  	towerIterator.RemoveCandidate(secondTower.ID, nil)
   138  	assertActiveCandidate(t, towerIterator, secondTower, false)
   139  	assertNextCandidate(t, towerIterator, thirdTower)
   140  
   141  	// We'll then update the fourth candidate with a new address. A
   142  	// duplicate shouldn't be added since it already exists within the
   143  	// iterator, but the new address should be.
   144  	fourthTower := towers[3]
   145  	assertActiveCandidate(t, towerIterator, fourthTower, true)
   146  	fourthTower.AddAddress(randAddr(t))
   147  	towerIterator.AddCandidate(fourthTower)
   148  	assertNextCandidate(t, towerIterator, fourthTower)
   149  
   150  	// Finally, we'll attempt to add a new candidate to the end of the
   151  	// iterator. Since it didn't already exist and we've reached the end, it
   152  	// should be available as the next candidate.
   153  	towerIterator.AddCandidate(secondTower)
   154  	assertActiveCandidate(t, towerIterator, secondTower, true)
   155  	assertNextCandidate(t, towerIterator, secondTower)
   156  }