dex: porting test to enode (#9)

5 files changed, 141 insertions, 112 deletions

diff --git a/dex/helper_test.go b/dex/helper_test.go
index a077304b1..9a135d3e3 100644
--- a/dex/helper_test.go
+++ b/dex/helper_test.go
@@ -21,8 +21,8 @@ package dex
 
 import (
 	"crypto/ecdsa"
-	"crypto/rand"
 	"math/big"
+	"net"
 	"sort"
 	"sync"
 	"testing"
@@ -48,17 +48,20 @@ var (
 
 // testP2PServer is a fake, helper p2p server for testing purposes.
 type testP2PServer struct {
-	mu     sync.Mutex
-	self   *enode.Node
-	direct map[enode.NodeID]*enode.Node
-	group  map[string][]*enode.Node
+	mu      sync.Mutex
+	self    *enode.Node
+	privkey *ecdsa.PrivateKey
+	direct  map[enode.ID]*enode.Node
+	group   map[string][]*enode.Node
 }
 
-func newTestP2PServer(self *enode.Node) *testP2PServer {
+func newTestP2PServer(privkey *ecdsa.PrivateKey) *testP2PServer {
+	self := enode.NewV4(&privkey.PublicKey, net.IP{}, 0, 0)
 	return &testP2PServer{
-		self:   self,
-		direct: make(map[enode.NodeID]*enode.Node),
-		group:  make(map[string][]*enode.Node),
+		self:    self,
+		privkey: privkey,
+		direct:  make(map[enode.ID]*enode.Node),
+		group:   make(map[string][]*enode.Node),
 	}
 }
 
@@ -66,16 +69,20 @@ func (s *testP2PServer) Self() *enode.Node {
 	return s.self
 }
 
+func (s *testP2PServer) GetPrivateKey() *ecdsa.PrivateKey {
+	return s.privkey
+}
+
 func (s *testP2PServer) AddDirectPeer(node *enode.Node) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	s.direct[node.ID] = node
+	s.direct[node.ID()] = node
 }
 
 func (s *testP2PServer) RemoveDirectPeer(node *enode.Node) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	delete(s.direct, node.ID)
+	delete(s.direct, node.ID())
 }
 
 func (s *testP2PServer) AddGroup(
@@ -104,7 +111,7 @@ func newTestProtocolManager(mode downloader.SyncMode, blocks int, generator func
 			Alloc:  core.GenesisAlloc{testBank: {Balance: big.NewInt(1000000), Staked: big.NewInt(0)}},
 		}
 		genesis       = gspec.MustCommit(db)
-		blockchain, _ = core.NewBlockChain(db, nil, gspec.Config, engine, vm.Config{})
+		blockchain, _ = core.NewBlockChain(db, nil, gspec.Config, engine, vm.Config{}, nil)
 	)
 	chain, _ := core.GenerateChain(gspec.Config, genesis, ethash.NewFaker(), db, blocks, generator)
 	if _, err := blockchain.InsertChain(chain); err != nil {
@@ -122,7 +129,11 @@ func newTestProtocolManager(mode downloader.SyncMode, blocks int, generator func
 	if err != nil {
 		return nil, nil, err
 	}
-	pm.Start(newTestP2PServer(&enode.Node{}), 1000)
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		return nil, nil, err
+	}
+	pm.Start(newTestP2PServer(key), 1000)
 	return pm, db, nil
 }
 
@@ -222,13 +233,16 @@ type testPeer struct {
 // newTestPeer creates a new peer registered at the given protocol manager.
 func newTestPeer(name string, version int, pm *ProtocolManager, shake bool) (*testPeer, <-chan error) {
 	// Create a message pipe to communicate through
-	app, net := p2p.MsgPipe()
+	app, pipenet := p2p.MsgPipe()
 
-	// Generate a random id and create the peer
-	var id enode.NodeID
-	rand.Read(id[:])
+	// Generate a random key and create the peer
+	key, err := crypto.GenerateKey()
+	if err != nil {
+		panic(err)
+	}
 
-	peer := pm.newPeer(version, p2p.NewPeer(id, name, nil), net)
+	node := enode.NewV4(&key.PublicKey, net.IP{}, 0, 0)
+	peer := pm.newPeer(version, p2p.NewPeerWithEnode(node, name, nil), pipenet)
 
 	// Start the peer on a new thread
 	errc := make(chan error, 1)
@@ -240,7 +254,7 @@ func newTestPeer(name string, version int, pm *ProtocolManager, shake bool) (*te
 			errc <- p2p.DiscQuitting
 		}
 	}()
-	tp := &testPeer{app: app, net: net, peer: peer}
+	tp := &testPeer{app: app, net: pipenet, peer: peer}
 	// Execute any implicitly requested handshakes and return
 	if shake {
 		var (
diff --git a/dex/nodetable_test.go b/dex/nodetable_test.go
index 2e44eabe5..5adf0fbec 100644
--- a/dex/nodetable_test.go
+++ b/dex/nodetable_test.go
@@ -2,12 +2,12 @@ package dex
 
 import (
 	"crypto/ecdsa"
+	"net"
 	"testing"
 	"time"
 
 	"github.com/dexon-foundation/dexon/common"
 	"github.com/dexon-foundation/dexon/crypto"
-	"github.com/dexon-foundation/dexon/p2p/discover"
 	"github.com/dexon-foundation/dexon/p2p/enode"
 )
 
@@ -87,7 +87,7 @@ func TestNodeTable(t *testing.T) {
 	}
 }
 
-func randomID() enode.ID {
+func randomEnode() *enode.Node {
 	var err error
 	var privkey *ecdsa.PrivateKey
 	for {
@@ -96,5 +96,9 @@ func randomID() enode.ID {
 			break
 		}
 	}
-	return discover.PubkeyID(&(privkey.PublicKey))
+	return enode.NewV4(&privkey.PublicKey, net.IP{}, 0, 0)
+}
+
+func randomID() enode.ID {
+	return randomEnode().ID()
 }
diff --git a/dex/peer_test.go b/dex/peer_test.go
index a4dc9269f..6c0804498 100644
--- a/dex/peer_test.go
+++ b/dex/peer_test.go
@@ -1,22 +1,21 @@
 package dex
 
 import (
+	"encoding/hex"
 	"fmt"
-	"math/big"
 	"testing"
 
 	"github.com/dexon-foundation/dexon/crypto"
-	"github.com/dexon-foundation/dexon/p2p/discover"
 	"github.com/dexon-foundation/dexon/p2p/enode"
 )
 
 func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
-	self := discover.Node{ID: nodeID(0)}
 	key, err := crypto.GenerateKey()
 	if err != nil {
 		t.Fatal(err)
 	}
-	server := newTestP2PServer(&self, key)
+	server := newTestP2PServer(key)
+	self := server.Self()
 	table := newNodeTable()
 
 	gov := &testGovernance{
@@ -25,25 +24,30 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		},
 	}
 
-	round10 := [][]enode.ID{
-		[]enode.ID{nodeID(0), nodeID(1), nodeID(2)},
-		[]enode.ID{nodeID(1), nodeID(3)},
-		[]enode.ID{nodeID(2), nodeID(4)},
+	var nodes []*enode.Node
+	for i := 0; i < 9; i++ {
+		nodes = append(nodes, randomEnode())
 	}
-	round11 := [][]enode.ID{
-		[]enode.ID{nodeID(0), nodeID(1), nodeID(5)},
-		[]enode.ID{nodeID(5), nodeID(6)},
-		[]enode.ID{nodeID(0), nodeID(2), nodeID(4)},
+
+	round10 := [][]*enode.Node{
+		[]*enode.Node{self, nodes[1], nodes[2]},
+		[]*enode.Node{nodes[1], nodes[3]},
+		[]*enode.Node{nodes[2], nodes[4]},
+	}
+	round11 := [][]*enode.Node{
+		[]*enode.Node{self, nodes[1], nodes[5]},
+		[]*enode.Node{nodes[5], nodes[6]},
+		[]*enode.Node{self, nodes[2], nodes[4]},
 	}
-	round12 := [][]enode.ID{
-		[]enode.ID{nodeID(0), nodeID(3), nodeID(5)},
-		[]enode.ID{nodeID(0), nodeID(7), nodeID(8)},
-		[]enode.ID{nodeID(0), nodeID(2), nodeID(6)},
+	round12 := [][]*enode.Node{
+		[]*enode.Node{self, nodes[3], nodes[5]},
+		[]*enode.Node{self, nodes[7], nodes[8]},
+		[]*enode.Node{self, nodes[2], nodes[6]},
 	}
 
 	gov.notarySetFunc = func(
 		round uint64, cid uint32) (map[string]struct{}, error) {
-		m := map[uint64][][]enode.ID{
+		m := map[uint64][][]*enode.Node{
 			10: round10,
 			11: round11,
 			12: round12,
@@ -52,8 +56,8 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 	}
 
 	ps := newPeerSet(gov, server, table)
-	peer1 := newDummyPeer(nodeID(1))
-	peer2 := newDummyPeer(nodeID(2))
+	peer1 := newDummyPeer(nodes[1])
+	peer2 := newDummyPeer(nodes[2])
 	err = ps.Register(peer1)
 	if err != nil {
 		t.Error(err)
@@ -67,10 +71,10 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 	ps.BuildNotaryConn(10)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(1).String(): []peerLabel{
+		nodes[1].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 10},
 		},
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 10},
 		},
 	})
@@ -82,7 +86,7 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(1), nodeID(2),
+		nodes[1].ID(), nodes[2].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -99,18 +103,18 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 	ps.BuildNotaryConn(11)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(1).String(): []peerLabel{
+		nodes[1].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 10},
 			peerLabel{notaryset, 0, 11},
 		},
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 10},
 			peerLabel{notaryset, 2, 11},
 		},
-		nodeID(4).String(): []peerLabel{
+		nodes[4].ID().String(): []peerLabel{
 			peerLabel{notaryset, 2, 11},
 		},
-		nodeID(5).String(): []peerLabel{
+		nodes[5].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 11},
 		},
 	})
@@ -122,7 +126,7 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(1), nodeID(2), nodeID(4), nodeID(5),
+		nodes[1].ID(), nodes[2].ID(), nodes[4].ID(), nodes[5].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -140,32 +144,32 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 	ps.BuildNotaryConn(12)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(1).String(): []peerLabel{
+		nodes[1].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 10},
 			peerLabel{notaryset, 0, 11},
 		},
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 10},
 			peerLabel{notaryset, 2, 11},
 			peerLabel{notaryset, 2, 12},
 		},
-		nodeID(3).String(): []peerLabel{
+		nodes[3].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 12},
 		},
-		nodeID(4).String(): []peerLabel{
+		nodes[4].ID().String(): []peerLabel{
 			peerLabel{notaryset, 2, 11},
 		},
-		nodeID(5).String(): []peerLabel{
+		nodes[5].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 11},
 			peerLabel{notaryset, 0, 12},
 		},
-		nodeID(6).String(): []peerLabel{
+		nodes[6].ID().String(): []peerLabel{
 			peerLabel{notaryset, 2, 12},
 		},
-		nodeID(7).String(): []peerLabel{
+		nodes[7].ID().String(): []peerLabel{
 			peerLabel{notaryset, 1, 12},
 		},
-		nodeID(8).String(): []peerLabel{
+		nodes[8].ID().String(): []peerLabel{
 			peerLabel{notaryset, 1, 12},
 		},
 	})
@@ -177,8 +181,8 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(1), nodeID(2), nodeID(3), nodeID(4),
-		nodeID(5), nodeID(6), nodeID(7), nodeID(8),
+		nodes[1].ID(), nodes[2].ID(), nodes[3].ID(), nodes[4].ID(),
+		nodes[5].ID(), nodes[6].ID(), nodes[7].ID(), nodes[8].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -196,22 +200,22 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 	ps.ForgetNotaryConn(11)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{notaryset, 2, 12},
 		},
-		nodeID(3).String(): []peerLabel{
+		nodes[3].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 12},
 		},
-		nodeID(5).String(): []peerLabel{
+		nodes[5].ID().String(): []peerLabel{
 			peerLabel{notaryset, 0, 12},
 		},
-		nodeID(6).String(): []peerLabel{
+		nodes[6].ID().String(): []peerLabel{
 			peerLabel{notaryset, 2, 12},
 		},
-		nodeID(7).String(): []peerLabel{
+		nodes[7].ID().String(): []peerLabel{
 			peerLabel{notaryset, 1, 12},
 		},
-		nodeID(8).String(): []peerLabel{
+		nodes[8].ID().String(): []peerLabel{
 			peerLabel{notaryset, 1, 12},
 		},
 	})
@@ -223,8 +227,8 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(2), nodeID(3),
-		nodeID(5), nodeID(6), nodeID(7), nodeID(8),
+		nodes[2].ID(), nodes[3].ID(),
+		nodes[5].ID(), nodes[6].ID(), nodes[7].ID(), nodes[8].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -256,28 +260,33 @@ func TestPeerSetBuildAndForgetNotaryConn(t *testing.T) {
 }
 
 func TestPeerSetBuildDKGConn(t *testing.T) {
-	self := discover.Node{ID: nodeID(0)}
 	key, err := crypto.GenerateKey()
 	if err != nil {
 		t.Fatal(err)
 	}
-	server := newTestP2PServer(&self, key)
+	server := newTestP2PServer(key)
+	self := server.Self()
 	table := newNodeTable()
 
+	var nodes []*enode.Node
+	for i := 0; i < 6; i++ {
+		nodes = append(nodes, randomEnode())
+	}
+
 	gov := &testGovernance{}
 
 	gov.dkgSetFunc = func(round uint64) (map[string]struct{}, error) {
-		m := map[uint64][]enode.ID{
-			10: []enode.ID{nodeID(0), nodeID(1), nodeID(2)},
-			11: []enode.ID{nodeID(1), nodeID(2), nodeID(5)},
-			12: []enode.ID{nodeID(0), nodeID(3), nodeID(5)},
+		m := map[uint64][]*enode.Node{
+			10: []*enode.Node{self, nodes[1], nodes[2]},
+			11: []*enode.Node{nodes[1], nodes[2], nodes[5]},
+			12: []*enode.Node{self, nodes[3], nodes[5]},
 		}
 		return newTestNodeSet(m[round]), nil
 	}
 
 	ps := newPeerSet(gov, server, table)
-	peer1 := newDummyPeer(nodeID(1))
-	peer2 := newDummyPeer(nodeID(2))
+	peer1 := newDummyPeer(nodes[1])
+	peer2 := newDummyPeer(nodes[2])
 	err = ps.Register(peer1)
 	if err != nil {
 		t.Error(err)
@@ -291,10 +300,10 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 	ps.BuildDKGConn(10)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(1).String(): []peerLabel{
+		nodes[1].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 10},
 		},
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 10},
 		},
 	})
@@ -306,7 +315,7 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(1), nodeID(2),
+		nodes[1].ID(), nodes[2].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -316,10 +325,10 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 	ps.BuildDKGConn(11)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(1).String(): []peerLabel{
+		nodes[1].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 10},
 		},
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 10},
 		},
 	})
@@ -331,7 +340,7 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(1), nodeID(2),
+		nodes[1].ID(), nodes[2].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -341,16 +350,16 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 	ps.BuildDKGConn(12)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(1).String(): []peerLabel{
+		nodes[1].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 10},
 		},
-		nodeID(2).String(): []peerLabel{
+		nodes[2].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 10},
 		},
-		nodeID(3).String(): []peerLabel{
+		nodes[3].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 12},
 		},
-		nodeID(5).String(): []peerLabel{
+		nodes[5].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 12},
 		},
 	})
@@ -362,7 +371,7 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(1), nodeID(2), nodeID(3), nodeID(5),
+		nodes[1].ID(), nodes[2].ID(), nodes[3].ID(), nodes[5].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -372,10 +381,10 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 	ps.ForgetDKGConn(11)
 
 	err = checkPeer2Labels(ps, map[string][]peerLabel{
-		nodeID(3).String(): []peerLabel{
+		nodes[3].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 12},
 		},
-		nodeID(5).String(): []peerLabel{
+		nodes[5].ID().String(): []peerLabel{
 			peerLabel{dkgset, 0, 12},
 		},
 	})
@@ -387,7 +396,7 @@ func TestPeerSetBuildDKGConn(t *testing.T) {
 		t.Error(err)
 	}
 	err = checkDirectPeer(server, []enode.ID{
-		nodeID(3), nodeID(5),
+		nodes[3].ID(), nodes[5].ID(),
 	})
 	if err != nil {
 		t.Error(err)
@@ -487,21 +496,15 @@ func checkGroup(srvr *testP2PServer, want []string) error {
 	return nil
 }
 
-func nodeID(n int64) enode.ID {
-	b := big.NewInt(n).Bytes()
-	var id enode.ID
-	copy(id[len(id)-len(b):], b)
-	return id
-}
-
-func newTestNodeSet(nodes []enode.ID) map[string]struct{} {
+func newTestNodeSet(nodes []*enode.Node) map[string]struct{} {
 	m := make(map[string]struct{})
 	for _, node := range nodes {
-		m[node.String()] = struct{}{}
+		b := crypto.FromECDSAPub(node.Pubkey())
+		m[hex.EncodeToString(b)] = struct{}{}
 	}
 	return m
 }
 
-func newDummyPeer(id enode.ID) *peer {
-	return &peer{id: id.String()}
+func newDummyPeer(node *enode.Node) *peer {
+	return &peer{id: node.ID().String()}
 }
diff --git a/dex/protocol_test.go b/dex/protocol_test.go
index 2e9ba3470..b0e89230a 100644
--- a/dex/protocol_test.go
+++ b/dex/protocol_test.go
@@ -17,6 +17,8 @@
 package dex
 
 import (
+	"crypto/ecdsa"
+	"encoding/hex"
 	"fmt"
 	"reflect"
 	"sync"
@@ -34,7 +36,6 @@ import (
 	"github.com/dexon-foundation/dexon/crypto"
 	"github.com/dexon-foundation/dexon/eth/downloader"
 	"github.com/dexon-foundation/dexon/p2p"
-	"github.com/dexon-foundation/dexon/p2p/discover"
 	"github.com/dexon-foundation/dexon/rlp"
 )
 
@@ -72,7 +73,7 @@ func testStatusMsgErrors(t *testing.T, protocol int) {
 		},
 		{
 			code: StatusMsg, data: statusData{uint32(protocol), 999, td, head.Hash(), genesis.Hash()},
-			wantError: errResp(ErrNetworkIdMismatch, "999 (!= 1)"),
+			wantError: errResp(ErrNetworkIdMismatch, "999 (!= 237)"),
 		},
 		{
 			code: StatusMsg, data: statusData{uint32(protocol), DefaultConfig.NetworkId, td, head.Hash(), common.Hash{3}},
@@ -237,7 +238,7 @@ func TestRecvNodeMetas(t *testing.T) {
 	defer p.close()
 
 	meta := NodeMeta{
-		ID: nodeID(1),
+		ID: randomID(),
 	}
 
 	ch := make(chan newMetasEvent)
@@ -265,8 +266,8 @@ func TestSendNodeMetas(t *testing.T) {
 	defer pm.Stop()
 
 	allmetas := make([]*NodeMeta, 100)
-	for nonce := range allmetas {
-		allmetas[nonce] = &NodeMeta{ID: nodeID(int64(nonce))}
+	for i := 0; i < len(allmetas); i++ {
+		allmetas[i] = &NodeMeta{ID: randomID()}
 	}
 
 	// Connect several peers. They should all receive the pending transactions.
@@ -548,7 +549,8 @@ func TestSendVote(t *testing.T) {
 	for i, tt := range testPeers {
 		p, _ := newTestPeer(fmt.Sprintf("peer #%d", i), dex64, pm, true)
 		if tt.label != nil {
-			pm.peers.addDirectPeer(p.id, *tt.label)
+			b := crypto.FromECDSAPub(p.Node().Pubkey())
+			pm.peers.addDirectPeer(hex.EncodeToString(b), *tt.label)
 		}
 		wg.Add(1)
 		go checkvote(p, tt.isReceiver)
@@ -558,17 +560,14 @@ func TestSendVote(t *testing.T) {
 	wg.Wait()
 }
 
-type mockPublicKey struct {
-	id enode.ID
-}
+type mockPublicKey ecdsa.PublicKey
 
 func (p *mockPublicKey) VerifySignature(hash coreCommon.Hash, signature coreCrypto.Signature) bool {
 	return true
 }
 
 func (p *mockPublicKey) Bytes() []byte {
-	b, _ := p.id.Pubkey()
-	return crypto.FromECDSAPub(b)
+	return crypto.FromECDSAPub((*ecdsa.PublicKey)(p))
 }
 
 func TestRecvDKGPrivateShare(t *testing.T) {
@@ -627,7 +626,7 @@ func TestSendDKGPrivateShare(t *testing.T) {
 		},
 	}
 
-	go pm.SendDKGPrivateShare(&mockPublicKey{p1.ID()}, &privateShare)
+	go pm.SendDKGPrivateShare((*mockPublicKey)(p1.Node().Pubkey()), &privateShare)
 	msg, err := p1.app.ReadMsg()
 	if err != nil {
 		t.Errorf("%v: read error: %v", p1.Peer, err)
diff --git a/p2p/peer.go b/p2p/peer.go
index 2c357fdc9..3d6ce9dcd 100644
--- a/p2p/peer.go
+++ b/p2p/peer.go
@@ -125,6 +125,15 @@ func NewPeer(id enode.ID, name string, caps []Cap) *Peer {
 	return peer
 }
 
+// NewPeer returns peer with given enode.Node for testing purposes.
+func NewPeerWithEnode(node *enode.Node, name string, caps []Cap) *Peer {
+	pipe, _ := net.Pipe()
+	conn := &conn{fd: pipe, transport: nil, node: node, caps: caps, name: name}
+	peer := newPeer(conn, nil)
+	close(peer.closed) // ensures Disconnect doesn't block
+	return peer
+}
+
 // ID returns the node's public key.
 func (p *Peer) ID() enode.ID {
 	return p.rw.node.ID()