watchtower/wtmock/peer: set local pubkey

This commit is contained in:
Conner Fromknecht 2019-03-15 02:34:00 -07:00
parent 80040d9d96
commit 05e3a7f6c0
No known key found for this signature in database
GPG Key ID: E7D737B67FA592C7
2 changed files with 18 additions and 9 deletions

View File

@ -27,14 +27,17 @@ type MockPeer struct {
}
// NewMockPeer returns a fresh MockPeer.
func NewMockPeer(pk *btcec.PublicKey, addr net.Addr, bufferSize int) *MockPeer {
func NewMockPeer(lpk, rpk *btcec.PublicKey, addr net.Addr,
bufferSize int) *MockPeer {
return &MockPeer{
remotePub: pk,
remotePub: rpk,
remoteAddr: addr,
localAddr: &net.TCPAddr{
IP: net.IP{0x32, 0x31, 0x30, 0x29},
Port: 36723,
},
localPub: lpk,
IncomingMsgs: make(chan []byte, bufferSize),
OutgoingMsgs: make(chan []byte, bufferSize),
Quit: make(chan struct{}),

View File

@ -87,10 +87,12 @@ func TestServerOnlyAcceptOnePeer(t *testing.T) {
s := initServer(t, nil, timeoutDuration)
defer s.Stop()
localPub := randPubKey(t)
// Create two peers using the same session id.
peerPub := randPubKey(t)
peer1 := wtmock.NewMockPeer(peerPub, nil, 0)
peer2 := wtmock.NewMockPeer(peerPub, nil, 0)
peer1 := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
peer2 := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
// Serialize a Init message to be sent by both peers.
init := wtwire.NewInitMessage(
@ -219,9 +221,11 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
s := initServer(t, nil, timeoutDuration)
defer s.Stop()
localPub := randPubKey(t)
// Create a new client and connect to server.
peerPub := randPubKey(t)
peer := wtmock.NewMockPeer(peerPub, nil, 0)
peer := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Send the CreateSession message, and wait for a reply.
@ -249,7 +253,7 @@ func testServerCreateSession(t *testing.T, i int, test createSessionTestCase) {
// Simulate a peer with the same session id connection to the server
// again.
peer = wtmock.NewMockPeer(peerPub, nil, 0)
peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Send the _same_ CreateSession message as the first attempt.
@ -559,9 +563,11 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
s := initServer(t, nil, timeoutDuration)
defer s.Stop()
localPub := randPubKey(t)
// Create a new client and connect to the server.
peerPub := randPubKey(t)
peer := wtmock.NewMockPeer(peerPub, nil, 0)
peer := wtmock.NewMockPeer(localPub, peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Register a session for this client to use in the subsequent tests.
@ -581,7 +587,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 = wtmock.NewMockPeer(peerPub, nil, 0)
peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
// Send the intended StateUpdate messages in series.
@ -592,7 +598,7 @@ func testServerStateUpdates(t *testing.T, i int, test stateUpdateTestCase) {
if update == nil {
assertConnClosed(t, peer, 2*timeoutDuration)
peer = wtmock.NewMockPeer(peerPub, nil, 0)
peer = wtmock.NewMockPeer(localPub, peerPub, nil, 0)
connect(t, i, s, peer, test.initMsg, timeoutDuration)
continue