lnd.xprv/discovery/mock_test.go

134 lines
2.9 KiB
Go
Raw Normal View History

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
}