github.com/decred/dcrlnd@v0.7.6/discovery/mock_test.go (about)

     1  package discovery
     2  
     3  import (
     4  	"errors"
     5  	"net"
     6  	"sync"
     7  
     8  	"github.com/decred/dcrd/dcrec/secp256k1/v4"
     9  	"github.com/decred/dcrd/wire"
    10  	"github.com/decred/dcrlnd/channeldb"
    11  	"github.com/decred/dcrlnd/lnpeer"
    12  	"github.com/decred/dcrlnd/lnwire"
    13  )
    14  
    15  // mockPeer implements the lnpeer.Peer interface and is used to test the
    16  // gossiper's interaction with peers.
    17  type mockPeer struct {
    18  	pk       *secp256k1.PublicKey
    19  	sentMsgs chan lnwire.Message
    20  	quit     chan struct{}
    21  }
    22  
    23  var _ lnpeer.Peer = (*mockPeer)(nil)
    24  
    25  func (p *mockPeer) SendMessage(_ bool, msgs ...lnwire.Message) error {
    26  	if p.sentMsgs == nil && p.quit == nil {
    27  		return nil
    28  	}
    29  
    30  	for _, msg := range msgs {
    31  		select {
    32  		case p.sentMsgs <- msg:
    33  		case <-p.quit:
    34  			return errors.New("peer disconnected")
    35  		}
    36  	}
    37  
    38  	return nil
    39  }
    40  
    41  func (p *mockPeer) SendMessageLazy(sync bool, msgs ...lnwire.Message) error {
    42  	return p.SendMessage(sync, msgs...)
    43  }
    44  
    45  func (p *mockPeer) AddNewChannel(_ *channeldb.OpenChannel, _ <-chan struct{}) error {
    46  	return nil
    47  }
    48  func (p *mockPeer) WipeChannel(_ *wire.OutPoint)      {}
    49  func (p *mockPeer) IdentityKey() *secp256k1.PublicKey { return p.pk }
    50  func (p *mockPeer) PubKey() [33]byte {
    51  	var pubkey [33]byte
    52  	copy(pubkey[:], p.pk.SerializeCompressed())
    53  	return pubkey
    54  }
    55  func (p *mockPeer) Address() net.Addr { return nil }
    56  func (p *mockPeer) QuitSignal() <-chan struct{} {
    57  	return p.quit
    58  }
    59  func (p *mockPeer) Inbound() bool { return false }
    60  func (p *mockPeer) LocalFeatures() *lnwire.FeatureVector {
    61  	return nil
    62  }
    63  func (p *mockPeer) RemoteFeatures() *lnwire.FeatureVector {
    64  	return nil
    65  }
    66  
    67  // mockMessageStore is an in-memory implementation of the MessageStore interface
    68  // used for the gossiper's unit tests.
    69  type mockMessageStore struct {
    70  	sync.Mutex
    71  	messages map[[33]byte]map[lnwire.Message]struct{}
    72  }
    73  
    74  func newMockMessageStore() *mockMessageStore {
    75  	return &mockMessageStore{
    76  		messages: make(map[[33]byte]map[lnwire.Message]struct{}),
    77  	}
    78  }
    79  
    80  var _ GossipMessageStore = (*mockMessageStore)(nil)
    81  
    82  func (s *mockMessageStore) AddMessage(msg lnwire.Message, pubKey [33]byte) error {
    83  	s.Lock()
    84  	defer s.Unlock()
    85  
    86  	if _, ok := s.messages[pubKey]; !ok {
    87  		s.messages[pubKey] = make(map[lnwire.Message]struct{})
    88  	}
    89  
    90  	s.messages[pubKey][msg] = struct{}{}
    91  
    92  	return nil
    93  }
    94  
    95  func (s *mockMessageStore) DeleteMessage(msg lnwire.Message, pubKey [33]byte) error {
    96  	s.Lock()
    97  	defer s.Unlock()
    98  
    99  	peerMsgs, ok := s.messages[pubKey]
   100  	if !ok {
   101  		return nil
   102  	}
   103  
   104  	delete(peerMsgs, msg)
   105  	return nil
   106  }
   107  
   108  func (s *mockMessageStore) Messages() (map[[33]byte][]lnwire.Message, error) {
   109  	s.Lock()
   110  	defer s.Unlock()
   111  
   112  	msgs := make(map[[33]byte][]lnwire.Message, len(s.messages))
   113  	for peer, peerMsgs := range s.messages {
   114  		for msg := range peerMsgs {
   115  			msgs[peer] = append(msgs[peer], msg)
   116  		}
   117  	}
   118  	return msgs, nil
   119  }
   120  
   121  func (s *mockMessageStore) Peers() (map[[33]byte]struct{}, error) {
   122  	s.Lock()
   123  	defer s.Unlock()
   124  
   125  	peers := make(map[[33]byte]struct{}, len(s.messages))
   126  	for peer := range s.messages {
   127  		peers[peer] = struct{}{}
   128  	}
   129  	return peers, nil
   130  }
   131  
   132  func (s *mockMessageStore) MessagesForPeer(pubKey [33]byte) ([]lnwire.Message, error) {
   133  	s.Lock()
   134  	defer s.Unlock()
   135  
   136  	peerMsgs, ok := s.messages[pubKey]
   137  	if !ok {
   138  		return nil, nil
   139  	}
   140  
   141  	msgs := make([]lnwire.Message, 0, len(peerMsgs))
   142  	for msg := range peerMsgs {
   143  		msgs = append(msgs, msg)
   144  	}
   145  
   146  	return msgs, nil
   147  }