server+discovery: replace gossiper message store with MessageStore

discovery/mock_test.go

@@ -0,0 +1,133 @@
+package discovery
+
+import (
+	"net"
+	"sync"
+
+	"github.com/btcsuite/btcd/btcec"
+	"github.com/btcsuite/btcd/wire"
+	"github.com/lightningnetwork/lnd/channeldb"
+	"github.com/lightningnetwork/lnd/lnpeer"
+	"github.com/lightningnetwork/lnd/lnwire"
+)
+
+// mockPeer implements the lnpeer.Peer interface and is used to test the
+// gossiper's interaction with peers.
+type mockPeer struct {
+	pk       *btcec.PublicKey
+	sentMsgs chan lnwire.Message
+	quit     chan struct{}
+}
+
+var _ lnpeer.Peer = (*mockPeer)(nil)
+
+func (p *mockPeer) SendMessage(_ bool, msgs ...lnwire.Message) error {
+	if p.sentMsgs == nil && p.quit == nil {
+		return nil
+	}
+
+	for _, msg := range msgs {
+		select {
+		case p.sentMsgs <- msg:
+		case <-p.quit:
+		}
+	}
+
+	return nil
+}
+func (p *mockPeer) AddNewChannel(_ *channeldb.OpenChannel, _ <-chan struct{}) error {
+	return nil
+}
+func (p *mockPeer) WipeChannel(_ *wire.OutPoint) error { return nil }
+func (p *mockPeer) IdentityKey() *btcec.PublicKey      { return p.pk }
+func (p *mockPeer) PubKey() [33]byte {
+	var pubkey [33]byte
+	copy(pubkey[:], p.pk.SerializeCompressed())
+	return pubkey
+}
+func (p *mockPeer) Address() net.Addr { return nil }
+func (p *mockPeer) QuitSignal() <-chan struct{} {
+	return p.quit
+}
+
+// mockMessageStore is an in-memory implementation of the MessageStore interface
+// used for the gossiper's unit tests.
+type mockMessageStore struct {
+	sync.Mutex
+	messages map[[33]byte]map[lnwire.Message]struct{}
+}
+
+func newMockMessageStore() *mockMessageStore {
+	return &mockMessageStore{
+		messages: make(map[[33]byte]map[lnwire.Message]struct{}),
+	}
+}
+
+var _ GossipMessageStore = (*mockMessageStore)(nil)
+
+func (s *mockMessageStore) AddMessage(msg lnwire.Message, pubKey [33]byte) error {
+	s.Lock()
+	defer s.Unlock()
+
+	if _, ok := s.messages[pubKey]; !ok {
+		s.messages[pubKey] = make(map[lnwire.Message]struct{})
+	}
+
+	s.messages[pubKey][msg] = struct{}{}
+
+	return nil
+}
+
+func (s *mockMessageStore) DeleteMessage(msg lnwire.Message, pubKey [33]byte) error {
+	s.Lock()
+	defer s.Unlock()
+
+	peerMsgs, ok := s.messages[pubKey]
+	if !ok {
+		return nil
+	}
+
+	delete(peerMsgs, msg)
+	return nil
+}
+
+func (s *mockMessageStore) Messages() (map[[33]byte][]lnwire.Message, error) {
+	s.Lock()
+	defer s.Unlock()
+
+	msgs := make(map[[33]byte][]lnwire.Message, len(s.messages))
+	for peer, peerMsgs := range s.messages {
+		for msg := range peerMsgs {
+			msgs[peer] = append(msgs[peer], msg)
+		}
+	}
+	return msgs, nil
+}
+
+func (s *mockMessageStore) Peers() (map[[33]byte]struct{}, error) {
+	s.Lock()
+	defer s.Unlock()
+
+	peers := make(map[[33]byte]struct{}, len(s.messages))
+	for peer := range s.messages {
+		peers[peer] = struct{}{}
+	}
+	return peers, nil
+}
+
+func (s *mockMessageStore) MessagesForPeer(pubKey [33]byte) ([]lnwire.Message, error) {
+	s.Lock()
+	defer s.Unlock()
+
+	peerMsgs, ok := s.messages[pubKey]
+	if !ok {
+		return nil, nil
+	}
+
+	msgs := make([]lnwire.Message, 0, len(peerMsgs))
+	for msg := range peerMsgs {
+		msgs = append(msgs, msg)
+	}
+
+	return msgs, nil
+}