From 37d5a2cf3ed9f1562bbab946e960bf54eb41198d Mon Sep 17 00:00:00 2001
From: Anton Kalyaev <anton.kalyaev@gmail.com>
Date: Wed, 11 Jan 2017 19:23:31 +0400
Subject: [PATCH] implement RemovePeer for PEXReactor

---
 pex_reactor.go       |  7 +++++--
 pex_reactor_tests.go | 50 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 pex_reactor_tests.go

diff --git a/pex_reactor.go b/pex_reactor.go
index 01d0cbd18..8d49de992 100644
--- a/pex_reactor.go
+++ b/pex_reactor.go
@@ -73,13 +73,16 @@ func (r *PEXReactor) AddPeer(p *Peer) {
 			r.RequestPEX(p)
 		}
 	} else { // For inbound connections, the peer is its own source
-		r.book.AddAddress(netAddr, netAddr)
+		addr := NewNetAddressString(p.ListenAddr)
+		r.book.AddAddress(addr, addr)
 	}
 }
 
 // RemovePeer implements Reactor
 func (r *PEXReactor) RemovePeer(p *Peer, reason interface{}) {
-	// TODO
+	addr := NewNetAddressString(p.ListenAddr)
+	// addr will be ejected from the book
+	r.book.MarkBad(addr)
 }
 
 // Receive implements Reactor by handling incoming PEX messages.
diff --git a/pex_reactor_tests.go b/pex_reactor_tests.go
new file mode 100644
index 000000000..5825495f7
--- /dev/null
+++ b/pex_reactor_tests.go
@@ -0,0 +1,50 @@
+package p2p
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	. "github.com/tendermint/go-common"
+)
+
+func TestBasic(t *testing.T) {
+	book := NewAddrBook(createTempFileName("addrbook"), true)
+	r := NewPEXReactor(book)
+
+	assert.NotNil(t, r)
+	assert.NotEmpty(t, r.GetChannels())
+}
+
+func TestAddRemovePeer(t *testing.T) {
+	book := NewAddrBook(createTempFileName("addrbook"), true)
+	r := NewPEXReactor(book)
+
+	size := book.Size()
+	peer := createRandomPeer(false)
+
+	r.AddPeer(peer)
+	assert.Equal(t, size+1, book.Size())
+
+	r.RemovePeer(peer, "peer not available")
+	assert.Equal(t, size, book.Size())
+
+	outboundPeer := createRandomPeer(true)
+
+	r.AddPeer(outboundPeer)
+	assert.Equal(t, size, book.Size(), "size must not change")
+
+	r.RemovePeer(outboundPeer, "peer not available")
+	assert.Equal(t, size, book.Size(), "size must not change")
+}
+
+func createRandomPeer(outbound bool) *Peer {
+	return &Peer{
+		Key: RandStr(12),
+		NodeInfo: &NodeInfo{
+			RemoteAddr: Fmt("%v.%v.%v.%v:46656", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256),
+			ListenAddr: Fmt("%v.%v.%v.%v:46656", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256),
+		},
+		outbound: outbound,
+	}
+}