1 files changed, 15 insertions, 32 deletions

diff --git a/swarm/pss/pss_test.go b/swarm/pss/pss_test.go
index 41b03db28..6ba04cb5d 100644
--- a/swarm/pss/pss_test.go
+++ b/swarm/pss/pss_test.go
@@ -556,23 +556,6 @@ OUTER:
 	}
 }
 
-type pssTestPeer struct {
-	*protocols.Peer
-	addr []byte
-}
-
-func (t *pssTestPeer) Address() []byte {
-	return t.addr
-}
-
-func (t *pssTestPeer) Update(addr network.OverlayAddr) network.OverlayAddr {
-	return addr
-}
-
-func (t *pssTestPeer) Off() network.OverlayAddr {
-	return &pssTestPeer{}
-}
-
 // forwarding should skip peers that do not have matching pss capabilities
 func TestMismatch(t *testing.T) {
 
@@ -582,7 +565,7 @@ func TestMismatch(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	// initialize overlay
+	// initialize kad
 	baseaddr := network.RandomAddr()
 	kad := network.NewKademlia((baseaddr).Over(), network.NewKadParams())
 	rw := &p2p.MsgPipeRW{}
@@ -594,10 +577,10 @@ func TestMismatch(t *testing.T) {
 		Version: 0,
 	}
 	nid, _ := discover.HexID("0x01")
-	wrongpsspeer := &pssTestPeer{
-		Peer: protocols.NewPeer(p2p.NewPeer(nid, common.ToHex(wrongpssaddr.Over()), []p2p.Cap{wrongpsscap}), rw, nil),
-		addr: wrongpssaddr.Over(),
-	}
+	wrongpsspeer := network.NewPeer(&network.BzzPeer{
+		Peer:    protocols.NewPeer(p2p.NewPeer(nid, common.ToHex(wrongpssaddr.Over()), []p2p.Cap{wrongpsscap}), rw, nil),
+		BzzAddr: &network.BzzAddr{OAddr: wrongpssaddr.Over(), UAddr: nil},
+	}, kad)
 
 	// one peer doesn't even have pss (boo!)
 	nopssaddr := network.RandomAddr()
@@ -606,16 +589,16 @@ func TestMismatch(t *testing.T) {
 		Version: 1,
 	}
 	nid, _ = discover.HexID("0x02")
-	nopsspeer := &pssTestPeer{
-		Peer: protocols.NewPeer(p2p.NewPeer(nid, common.ToHex(nopssaddr.Over()), []p2p.Cap{nopsscap}), rw, nil),
-		addr: nopssaddr.Over(),
-	}
+	nopsspeer := network.NewPeer(&network.BzzPeer{
+		Peer:    protocols.NewPeer(p2p.NewPeer(nid, common.ToHex(nopssaddr.Over()), []p2p.Cap{nopsscap}), rw, nil),
+		BzzAddr: &network.BzzAddr{OAddr: nopssaddr.Over(), UAddr: nil},
+	}, kad)
 
 	// add peers to kademlia and activate them
 	// it's safe so don't check errors
-	kad.Register([]network.OverlayAddr{wrongpsspeer})
+	kad.Register(wrongpsspeer.BzzAddr)
 	kad.On(wrongpsspeer)
-	kad.Register([]network.OverlayAddr{nopsspeer})
+	kad.Register(nopsspeer.BzzAddr)
 	kad.On(nopsspeer)
 
 	// create pss
@@ -1636,17 +1619,17 @@ func newServices(allowRaw bool) adapters.Services {
 	}
 }
 
-func newTestPss(privkey *ecdsa.PrivateKey, overlay network.Overlay, ppextra *PssParams) *Pss {
+func newTestPss(privkey *ecdsa.PrivateKey, kad *network.Kademlia, ppextra *PssParams) *Pss {
 
 	var nid discover.NodeID
 	copy(nid[:], crypto.FromECDSAPub(&privkey.PublicKey))
 	addr := network.NewAddrFromNodeID(nid)
 
 	// set up routing if kademlia is not passed to us
-	if overlay == nil {
+	if kad == nil {
 		kp := network.NewKadParams()
 		kp.MinProxBinSize = 3
-		overlay = network.NewKademlia(addr.Over(), kp)
+		kad = network.NewKademlia(addr.Over(), kp)
 	}
 
 	// create pss
@@ -1654,7 +1637,7 @@ func newTestPss(privkey *ecdsa.PrivateKey, overlay network.Overlay, ppextra *Pss
 	if ppextra != nil {
 		pp.SymKeyCacheCapacity = ppextra.SymKeyCacheCapacity
 	}
-	ps, err := NewPss(overlay, pp)
+	ps, err := NewPss(kad, pp)
 	if err != nil {
 		return nil
 	}