lnd/discovery/mock_test.go
Conner Fromknecht ec784db511
multi: remove returned error from WipeChannel
The linter complains about not checking the return value from
WipeChannel in certain places. Instead of checking we simply remove the
returned error because the in-memory modifications cannot fail.
2020-04-02 17:39:29 -07:00

147 lines
3.2 KiB
Go

package discovery
import (
"errors"
"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 errors.New("peer disconnected")
}
}
return nil
}
func (p *mockPeer) SendMessageLazy(sync bool, msgs ...lnwire.Message) error {
return p.SendMessage(sync, msgs...)
}
func (p *mockPeer) AddNewChannel(_ *channeldb.OpenChannel, _ <-chan struct{}) error {
return nil
}
func (p *mockPeer) WipeChannel(_ *wire.OutPoint) {}
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
}
func (p *mockPeer) LocalFeatures() *lnwire.FeatureVector {
return nil
}
func (p *mockPeer) RemoteFeatures() *lnwire.FeatureVector {
return nil
}
// 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
}