Merge pull request #2514 from cfromknecht/add-wtmock-pkg

watchtower/multi: move MockSigner+MockPeer to wtmock
This commit is contained in:
Olaoluwa Osuntokun
2019-01-31 19:13:47 -08:00
committed by GitHub
4 changed files with 122 additions and 69 deletions

View File

@@ -1,98 +0,0 @@
// +build dev
package wtserver
import (
"fmt"
"net"
"time"
"github.com/btcsuite/btcd/btcec"
)
type MockPeer struct {
remotePub *btcec.PublicKey
remoteAddr net.Addr
IncomingMsgs chan []byte
OutgoingMsgs chan []byte
writeDeadline <-chan time.Time
readDeadline <-chan time.Time
Quit chan struct{}
}
func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer {
return &MockPeer{
remotePub: pk,
remoteAddr: addr,
IncomingMsgs: make(chan []byte, bufferSize),
OutgoingMsgs: make(chan []byte, bufferSize),
Quit: make(chan struct{}),
}
}
func (p *MockPeer) Write(b []byte) (n int, err error) {
select {
case p.OutgoingMsgs <- b:
return len(b), nil
case <-p.writeDeadline:
return 0, fmt.Errorf("write timeout expired")
case <-p.Quit:
return 0, fmt.Errorf("connection closed")
}
}
func (p *MockPeer) Close() error {
select {
case <-p.Quit:
return fmt.Errorf("connection already closed")
default:
close(p.Quit)
return nil
}
}
func (p *MockPeer) ReadNextMessage() ([]byte, error) {
select {
case b := <-p.IncomingMsgs:
return b, nil
case <-p.readDeadline:
return nil, fmt.Errorf("read timeout expired")
case <-p.Quit:
return nil, fmt.Errorf("connection closed")
}
}
func (p *MockPeer) SetWriteDeadline(t time.Time) error {
if t.IsZero() {
p.writeDeadline = nil
return nil
}
duration := time.Until(t)
p.writeDeadline = time.After(duration)
return nil
}
func (p *MockPeer) SetReadDeadline(t time.Time) error {
if t.IsZero() {
p.readDeadline = nil
return nil
}
duration := time.Until(t)
p.readDeadline = time.After(duration)
return nil
}
func (p *MockPeer) RemotePub() *btcec.PublicKey {
return p.remotePub
}
func (p *MockPeer) RemoteAddr() net.Addr {
return p.remoteAddr
}

View File

@@ -15,6 +15,7 @@ import (
"github.com/lightningnetwork/lnd/lnwire"
"github.com/lightningnetwork/lnd/watchtower/blob"
"github.com/lightningnetwork/lnd/watchtower/wtdb"
"github.com/lightningnetwork/lnd/watchtower/wtmock"
"github.com/lightningnetwork/lnd/watchtower/wtserver"
"github.com/lightningnetwork/lnd/watchtower/wtwire"
)
@@ -85,8 +86,8 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
// Create two peers using the same session id.
peerPub := randPubKey(t)
peer1 := wtserver.NewMockPeer(peerPub, nil, 0)
peer2 := wtserver.NewMockPeer(peerPub, nil, 0)
peer1 := wtmock.NewMockPeer(peerPub, nil, 0)
peer2 := wtmock.NewMockPeer(peerPub, nil, 0)
// Serialize a Init message to be sent by both peers.
init := wtwire.NewInitMessage(
@@ -112,8 +113,8 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
// Try to send a message on either peer, and record the opposite peer as
// the one we assume to be rejected.
var (
rejectedPeer *wtserver.MockPeer
acceptedPeer *wtserver.MockPeer
rejectedPeer *wtmock.MockPeer
acceptedPeer *wtmock.MockPeer
)
select {
case peer1.IncomingMsgs <- msg:
@@ -217,7 +218,7 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
// Create a new client and connect to server.
peerPub := randPubKey(t)
peer := wtserver.NewMockPeer(peerPub, nil, 0)
peer := wtmock.NewMockPeer(peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Send the CreateSession message, and wait for a reply.
@@ -245,7 +246,7 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
// Simulate a peer with the same session id connection to the server
// again.
peer = wtserver.NewMockPeer(peerPub, nil, 0)
peer = wtmock.NewMockPeer(peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Send the _same_ CreateSession message as the first attempt.
@@ -525,7 +526,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
// Create a new client and connect to the server.
peerPub := randPubKey(t)
peer := wtserver.NewMockPeer(peerPub, nil, 0)
peer := wtmock.NewMockPeer(peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Register a session for this client to use in the subsequent tests.
@@ -545,7 +546,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
// Now that the original connection has been closed, connect a new
// client with the same session id.
peer = wtserver.NewMockPeer(peerPub, nil, 0)
peer = wtmock.NewMockPeer(peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Send the intended StateUpdate messages in series.
@@ -556,7 +557,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
if update == nil {
assertConnClosed(t, peer, 2*timeoutDuration)
peer = wtserver.NewMockPeer(peerPub, nil, 0)
peer = wtmock.NewMockPeer(peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
continue
@@ -580,7 +581,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
assertConnClosed(t, peer, 2*timeoutDuration)
}
func connect(t *testing.T, i int, s wtserver.Interface, peer *wtserver.MockPeer,
func connect(t *testing.T, i int, s wtserver.Interface, peer *wtmock.MockPeer,
initMsg *wtwire.Init, timeout time.Duration) {
s.InboundPeerConnected(peer)
@@ -588,9 +589,9 @@ func connect(t *testing.T, i int, s wtserver.Interface, peer *wtserver.MockPeer,
recvReply(t, i, "Init", peer, timeout)
}
// sendMsg sends a wtwire.Message message via a wtserver.MockPeer.
// sendMsg sends a wtwire.Message message via a wtmock.MockPeer.
func sendMsg(t *testing.T, i int, msg wtwire.Message,
peer *wtserver.MockPeer, timeout time.Duration) {
peer *wtmock.MockPeer, timeout time.Duration) {
t.Helper()
@@ -612,7 +613,7 @@ func sendMsg(t *testing.T, i int, msg wtwire.Message,
// expected reply type. The supported replies are CreateSessionReply and
// StateUpdateReply.
func recvReply(t *testing.T, i int, name string,
peer *wtserver.MockPeer, timeout time.Duration) wtwire.Message {
peer *wtmock.MockPeer, timeout time.Duration) wtwire.Message {
t.Helper()
@@ -656,7 +657,7 @@ func recvReply(t *testing.T, i int, name string,
// assertConnClosed checks that the peer's connection is closed before the
// timeout expires.
func assertConnClosed(t *testing.T, peer *wtserver.MockPeer, duration time.Duration) {
func assertConnClosed(t *testing.T, peer *wtmock.MockPeer, duration time.Duration) {
t.Helper()
select {